🪚 OmniGraph™ OmniPoint transformation utility (#99)

LayerZero-Labs · Dec 12, 2023 · 76e4c9a · 76e4c9a
1 parent 0656295
commit 76e4c9a
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 114 deletions.
diff --git a/packages/utils-evm-hardhat/src/omnigraph/builder.ts b/packages/utils-evm-hardhat/src/omnigraph/builder.ts
@@ -1,7 +1,7 @@
-import type { OmniGraphHardhat } from './types'
+import type { OmniGraphHardhat, OmniGraphHardhatTransformer } from './types'
 import { OmniGraphBuilder } from '@layerzerolabs/utils'
 import assert from 'assert'
-import { OmniGraphHardhatTransformer, createOmniGraphHardhatTransformer } from './transformations'
+import { createOmniGraphHardhatTransformer } from './transformations'
 
 /**
  * OmniGraphBuilderHardhat houses all hardhat-specific utilities for building OmniGraphs

diff --git a/packages/utils-evm-hardhat/src/omnigraph/transformations.ts b/packages/utils-evm-hardhat/src/omnigraph/transformations.ts
@@ -1,33 +1,57 @@
-import type { OmniEdge, OmniGraph, OmniNode } from '@layerzerolabs/utils'
+import type { OmniEdge, OmniNode } from '@layerzerolabs/utils'
 import { isOmniPoint } from '@layerzerolabs/utils'
 import { omniContractToPoint } from '@layerzerolabs/utils-evm'
 import { createContractFactory } from './coordinates'
-import type { OmniContractFactoryHardhat, OmniEdgeHardhat, OmniGraphHardhat, OmniNodeHardhat } from './types'
+import type {
+    OmniContractFactoryHardhat,
+    OmniEdgeHardhat,
+    OmniGraphHardhatTransformer,
+    OmniNodeHardhat,
+    OmniPointHardhatTransformer,
+} from './types'
 
-export const createOmniNodeHardhatTransformer =
-    (contractFactory: OmniContractFactoryHardhat = createContractFactory()) =>
-    async <TNodeConfig>({ contract, config }: OmniNodeHardhat<TNodeConfig>): Promise<OmniNode<TNodeConfig>> => {
-        const point = isOmniPoint(contract) ? contract : omniContractToPoint(await contractFactory(contract))
+/**
+ * Create a function capable of transforming `OmniPointHardhat` to a regular `OmniPoint`
+ * with some help from an `OmniContractFactoryHardhat`
+ *
+ * @param {OmniContractFactoryHardhat} contractFactory
+ * @returns {OmniPointHardhatTransformer}
+ */
+export const createOmniPointHardhatTransformer =
+    (contractFactory: OmniContractFactoryHardhat = createContractFactory()): OmniPointHardhatTransformer =>
+    async (point) =>
+        isOmniPoint(point) ? point : omniContractToPoint(await contractFactory(point))
 
-        return { point, config: config as TNodeConfig }
-    }
+/**
+ * Create a function capable of transforming `OmniNodeHardhat` to a regular `OmniNode`
+ * with some help from an `OmniPointHardhatTransformer`
+ *
+ * @param {OmniPointHardhatTransformer} [pointTransformer]
+ * @returns
+ */
+export const createOmniNodeHardhatTransformer =
+    (pointTransformer = createOmniPointHardhatTransformer()) =>
+    async <TNodeConfig>({ contract, config }: OmniNodeHardhat<TNodeConfig>): Promise<OmniNode<TNodeConfig>> => ({
+        point: await pointTransformer(contract),
+        config: config as TNodeConfig,
+    })
 
+/**
+ * Create a function capable of transforming `OmniEdgeHardhat` to a regular `OmniEdge`
+ * with some help from an `OmniPointHardhatTransformer`
+ *
+ * @param {OmniPointHardhatTransformer} [pointTransformer]
+ * @returns
+ */
 export const createOmniEdgeHardhatTransformer =
-    (contractFactory: OmniContractFactoryHardhat = createContractFactory()) =>
-    async <TEdgeConfig>({
-        from: fromContract,
-        to: toContract,
-        config,
-    }: OmniEdgeHardhat<TEdgeConfig>): Promise<OmniEdge<TEdgeConfig>> => {
-        const from = isOmniPoint(fromContract) ? fromContract : omniContractToPoint(await contractFactory(fromContract))
-        const to = isOmniPoint(toContract) ? toContract : omniContractToPoint(await contractFactory(toContract))
-
-        return { vector: { from, to }, config: config as TEdgeConfig }
-    }
-
-export type OmniGraphHardhatTransformer<TNodeConfig = unknown, TEdgeConfig = unknown> = (
-    graph: OmniGraphHardhat<TNodeConfig, TEdgeConfig>
-) => Promise<OmniGraph<TNodeConfig, TEdgeConfig>>
+    (pointTransformer = createOmniPointHardhatTransformer()) =>
+    async <TEdgeConfig>({ from, to, config }: OmniEdgeHardhat<TEdgeConfig>): Promise<OmniEdge<TEdgeConfig>> => ({
+        vector: {
+            from: await pointTransformer(from),
+            to: await pointTransformer(to),
+        },
+        config: config as TEdgeConfig,
+    })
 
 export const createOmniGraphHardhatTransformer =
     <TNodeConfig, TEdgeConfig>(

diff --git a/packages/utils-evm-hardhat/src/omnigraph/types.ts b/packages/utils-evm-hardhat/src/omnigraph/types.ts
@@ -1,4 +1,4 @@
-import type { OmniPoint, WithEid, WithOptionals } from '@layerzerolabs/utils'
+import type { OmniGraph, OmniPoint, WithEid, WithOptionals } from '@layerzerolabs/utils'
 import type { OmniContract } from '@layerzerolabs/utils-evm'
 import type { Deployment } from 'hardhat-deploy/dist/types'
 
@@ -50,3 +50,9 @@ export interface OmniGraphHardhat<TNodeConfig = unknown, TEdgeConfig = unknown>
 }
 
 export type OmniContractFactoryHardhat = (point: OmniPointHardhat) => OmniContract | Promise<OmniContract>
+
+export type OmniPointHardhatTransformer = (point: OmniPointHardhat | OmniPoint) => Promise<OmniPoint>
+
+export type OmniGraphHardhatTransformer<TNodeConfig = unknown, TEdgeConfig = unknown> = (
+    graph: OmniGraphHardhat<TNodeConfig, TEdgeConfig>
+) => Promise<OmniGraph<TNodeConfig, TEdgeConfig>>
diff --git a/packages/utils-evm-hardhat/test/omnigraph/transformations.test.ts b/packages/utils-evm-hardhat/test/omnigraph/transformations.test.ts
@@ -4,6 +4,7 @@ import {
     createOmniEdgeHardhatTransformer,
     createOmniGraphHardhatTransformer,
     createOmniNodeHardhatTransformer,
+    createOmniPointHardhatTransformer,
 } from '@/omnigraph/transformations'
 import { Contract } from '@ethersproject/contracts'
 import { endpointArbitrary, evmAddressArbitrary, nullableArbitrary, pointArbitrary } from '@layerzerolabs/test-utils'
@@ -16,140 +17,93 @@ describe('omnigraph/transformations', () => {
         address: nullableArbitrary(evmAddressArbitrary),
     })
 
-    describe('createOmniNodeHardhatTransformer', () => {
+    describe('createOmniPointHardhatTransformer', () => {
         it('should pass the original value if contract is already an OmniPoint', async () => {
             await fc.assert(
-                fc.asyncProperty(pointArbitrary, fc.anything(), async (point, config) => {
+                fc.asyncProperty(pointArbitrary, async (point) => {
                     const contractFactory = jest.fn().mockRejectedValue('Oh no')
-                    const transformer = createOmniNodeHardhatTransformer(contractFactory)
+                    const transformer = createOmniPointHardhatTransformer(contractFactory)
 
-                    const node = await transformer({ contract: point, config })
+                    const transformed = await transformer(point)
 
-                    expect(node).toEqual({ point, config })
+                    expect(transformed).toBe(point)
                     expect(contractFactory).not.toHaveBeenCalled()
                 })
             )
         })
 
         it('should call the contractFactory if contract is not an OmniPoint', async () => {
             await fc.assert(
-                fc.asyncProperty(
-                    pointHardhatArbitrary,
-                    evmAddressArbitrary,
-                    fc.anything(),
-                    async (point, address, config) => {
-                        fc.pre(!isOmniPoint(point))
-
-                        const contract = new Contract(address, [])
-                        const contractFactory = jest
-                            .fn()
-                            .mockImplementation(async (point: OmniPointHardhat) => ({ eid: point.eid, contract }))
-                        const transformer = createOmniNodeHardhatTransformer(contractFactory)
-
-                        const node = await transformer({ contract: point, config })
-
-                        expect(node).toEqual({ point: { eid: point.eid, address }, config })
-                        expect(contractFactory).toHaveBeenCalledTimes(1)
-                        expect(contractFactory).toHaveBeenCalledWith(point)
-                    }
-                )
-            )
-        })
-    })
+                fc.asyncProperty(pointHardhatArbitrary, evmAddressArbitrary, async (point, address) => {
+                    fc.pre(!isOmniPoint(point))
 
-    describe('createOmniEdgeHardhatTransformer', () => {
-        it('should pass the original values if from and to are already OmniPoints', async () => {
-            await fc.assert(
-                fc.asyncProperty(pointArbitrary, pointArbitrary, fc.anything(), async (from, to, config) => {
-                    const contractFactory = jest.fn().mockRejectedValue('Oh no')
-                    const transformer = createOmniEdgeHardhatTransformer(contractFactory)
+                    const contract = new Contract(address, [])
+                    const contractFactory = jest
+                        .fn()
+                        .mockImplementation(async (point: OmniPointHardhat) => ({ eid: point.eid, contract }))
+                    const transformer = createOmniPointHardhatTransformer(contractFactory)
 
-                    const edge = await transformer({ from, to, config })
+                    const transformed = await transformer(point)
 
-                    expect(edge).toEqual({ vector: { from, to }, config })
-                    expect(contractFactory).not.toHaveBeenCalled()
+                    expect(transformed).toEqual({ eid: point.eid, address })
+                    expect(contractFactory).toHaveBeenCalledTimes(1)
+                    expect(contractFactory).toHaveBeenCalledWith(point)
                 })
             )
         })
+    })
 
-        it('should call the contractFactory if from is not an OmniPoint', async () => {
-            await fc.assert(
-                fc.asyncProperty(
-                    pointHardhatArbitrary,
-                    pointArbitrary,
-                    evmAddressArbitrary,
-                    fc.anything(),
-                    async (from, to, address, config) => {
-                        fc.pre(!isOmniPoint(from))
-
-                        const contract = new Contract(address, [])
-                        const contractFactory = jest
-                            .fn()
-                            .mockImplementation(async (point: OmniPointHardhat) => ({ eid: point.eid, contract }))
-                        const transformer = createOmniEdgeHardhatTransformer(contractFactory)
-
-                        const edge = await transformer({ from, to, config })
-
-                        expect(edge).toEqual({ vector: { from: { eid: from.eid, address }, to }, config })
-                        expect(contractFactory).toHaveBeenCalledTimes(1)
-                        expect(contractFactory).toHaveBeenCalledWith(from)
-                    }
-                )
-            )
-        })
-
-        it('should call the contractFactory if to is not an OmniPoint', async () => {
+    describe('createOmniNodeHardhatTransformer', () => {
+        it('should call the pointTransformer on the point', async () => {
             await fc.assert(
                 fc.asyncProperty(
-                    pointArbitrary,
                     pointHardhatArbitrary,
                     evmAddressArbitrary,
                     fc.anything(),
-                    async (from, to, address, config) => {
-                        fc.pre(!isOmniPoint(to))
+                    async (point, address, config) => {
+                        fc.pre(!isOmniPoint(point))
 
-                        const contract = new Contract(address, [])
-                        const contractFactory = jest
-                            .fn()
-                            .mockImplementation(async (point: OmniPointHardhat) => ({ eid: point.eid, contract }))
-                        const transformer = createOmniEdgeHardhatTransformer(contractFactory)
+                        const pointTransformer = jest.fn().mockImplementation(async (point: OmniPointHardhat) => ({
+                            eid: point.eid,
+                            address: address,
+                        }))
+                        const transformer = createOmniNodeHardhatTransformer(pointTransformer)
 
-                        const edge = await transformer({ from, to, config })
+                        const node = await transformer({ contract: point, config })
 
-                        expect(edge).toEqual({ vector: { from, to: { eid: to.eid, address } }, config })
-                        expect(contractFactory).toHaveBeenCalledTimes(1)
-                        expect(contractFactory).toHaveBeenCalledWith(to)
+                        expect(node).toEqual({ point: { eid: point.eid, address }, config })
+                        expect(pointTransformer).toHaveBeenCalledTimes(1)
+                        expect(pointTransformer).toHaveBeenCalledWith(point)
                     }
                 )
             )
         })
+    })
 
-        it('should call the contractFactory if from & to are not OmniPoints', async () => {
+    describe('createOmniEdgeHardhatTransformer', () => {
+        it('should call the pointTransformer on from and to', async () => {
             await fc.assert(
                 fc.asyncProperty(
                     pointHardhatArbitrary,
                     pointHardhatArbitrary,
                     evmAddressArbitrary,
                     fc.anything(),
                     async (from, to, address, config) => {
-                        fc.pre(!isOmniPoint(from))
-                        fc.pre(!isOmniPoint(to))
-
-                        const contract = new Contract(address, [])
-                        const contractFactory = jest
-                            .fn()
-                            .mockImplementation(async (point: OmniPointHardhat) => ({ eid: point.eid, contract }))
-                        const transformer = createOmniEdgeHardhatTransformer(contractFactory)
+                        const pointTransformer = jest.fn().mockImplementation(async (point: OmniPointHardhat) => ({
+                            eid: point.eid,
+                            address,
+                        }))
+                        const transformer = createOmniEdgeHardhatTransformer(pointTransformer)
 
                         const edge = await transformer({ from, to, config })
 
                         expect(edge).toEqual({
                             vector: { from: { eid: from.eid, address }, to: { eid: to.eid, address } },
                             config,
                         })
-                        expect(contractFactory).toHaveBeenCalledTimes(2)
-                        expect(contractFactory).toHaveBeenCalledWith(from)
-                        expect(contractFactory).toHaveBeenCalledWith(to)
+                        expect(pointTransformer).toHaveBeenCalledTimes(2)
+                        expect(pointTransformer).toHaveBeenCalledWith(from)
+                        expect(pointTransformer).toHaveBeenCalledWith(to)
                     }
                 )
             )