Open-Attestation · osslgtm · Apr 26, 2023
diff --git a/src/implementations/testsHelper.ts b/src/implementations/testsHelper.ts
@@ -0,0 +1,138 @@
+import { BaseContract, BigNumber, constants, providers } from "ethers";
+
+export const AddressZero = constants.AddressZero;
+type defaultMockType = jest.Mock<any, any>;
+type SmartContractDataTypes = BigNumber | number | boolean | string; // bytes and unused not included
+type EthersResponseType = SmartContractDataTypes | Error;
+
+export const mockResolve = (value?: EthersResponseType): defaultMockType => {
+  const fn = jest.fn();
+  if (value instanceof Error) {
+    fn.mockRejectedValue(value);
+  } else {
+    fn.mockResolvedValue(value);
+  }
+  return fn;
+};
+
+export interface ValidContractMockParameters {
+  supportInterfaceValue?: boolean | Error;
+}
+
+interface MockContractInterface {
+  supportInterface: jest.Mock;
+  callStatic: {
+    supportInterface: jest.Mock;
+  };
+}
+
+export const getMockContract = ({
+  supportInterfaceValue = true,
+}: ValidContractMockParameters): MockContractInterface => {
+  const supportInterface = mockResolve(supportInterfaceValue);
+  return {
+    supportInterface,
+    callStatic: {
+      supportInterface,
+    },
+  };
+};
+
+export interface TokenRegistryMockParameters extends ValidContractMockParameters {
+  ownerOfValue?: string | Error;
+  address?: string;
+  titleEscrowFactoryAddress?: string;
+}
+
+interface MockTokenRegistryInterface {
+  ownerOf: jest.Mock;
+  genesis: jest.Mock;
+  titleEscrowFactory: jest.Mock;
+  supportInterfaces: jest.Mock;
+  callStatic: {
+    ownerOf: jest.Mock;
+    genesis: jest.Mock;
+    titleEscrowFactory: jest.Mock;
+    supportInterfaces: jest.Mock;
+  };
+}
+
+export const getMockTokenRegistry = ({
+  ownerOfValue = AddressZero,
+  supportInterfaceValue = true,
+  address = AddressZero,
+  titleEscrowFactoryAddress = AddressZero,
+}: TokenRegistryMockParameters): MockTokenRegistryInterface => {
+  const validContract = getMockContract({ supportInterfaceValue });
+  const ownerOf = mockResolve(ownerOfValue);
+  const genesis = mockResolve(BigNumber.from(0));
+  const titleEscrowFactory = mockResolve(titleEscrowFactoryAddress);
+  const contractFunctions = {
+    ownerOf,
+    genesis,
+    titleEscrowFactory,
+  };
+  const mockTokenRegistry = {
+    ...contractFunctions,
+    address: address,
+    callStatic: contractFunctions,
+  };
+  return mergeMockSmartContract({ base: validContract, override: mockTokenRegistry });
+};
+
+export interface TokenRegistryMockParameters extends ValidContractMockParameters {
+  getAddressValue?: string | Error;
+}
+
+export const initMockGetCode = (fn?: jest.Mock): void => {
+  if (!fn) {
+    const fn = jest.fn();
+    fn.mockResolvedValue(`0x`);
+  }
+  jest.spyOn(providers.BaseProvider.prototype, "getCode").mockImplementation(fn);
+};
+export interface WalletMockParameters {
+  codeValue?: string | Error;
+}
+
+export const getValidWalletContract = ({
+  codeValue = `0x`,
+}: WalletMockParameters): { provider: { getCode: jest.Mock } } => {
+  const getCode = mockResolve(codeValue);
+  return {
+    provider: {
+      getCode,
+    },
+  };
+};
+
+export interface MergeObjectParameters {
+  base: any;
+  override: any;
+}
+
+export const mergeMockSmartContract = ({ base, override }: MergeObjectParameters): any => {
+  override = mergeMockBaseContract(base, override, "functions");
+  override = mergeMockBaseContract(base, override, "callStatic");
+  override = mergeMockBaseContract(base, override, "estimateGas");
+  override = mergeMockBaseContract(base, override, "populateTransaction");
+  override = mergeMockBaseContract(base, override, "filters");
+  override = mergeMockBaseContract(base, override, "_runningEvents");
+  override = mergeMockBaseContract(base, override, "_wrappedEmits");
+  return {
+    ...base,
+    ...override,
+  };
+};
+
+const mergeMockBaseContract = <key extends keyof BaseContract>(base: any, override: any, keyName: key): any => {
+  if (keyName in override && keyName in base) {
+    if (typeof base[keyName] === "object" && typeof override[keyName] === "object") {
+      override[keyName] = {
+        ...base[keyName],
+        ...override[keyName],
+      };
+    }
+  }
+  return override;
+};
diff --git a/src/implementations/title-escrow/acceptSurrendered.test.ts b/src/implementations/title-escrow/acceptSurrendered.test.ts
@@ -2,13 +2,14 @@ import { TradeTrustToken__factory } from "@govtechsg/token-registry/contracts";
 import { Wallet } from "ethers";
 
 import { BaseTitleEscrowCommand as TitleEscrowSurrenderDocumentCommand } from "../../commands/title-escrow/title-escrow-command.type";
+import { AddressZero, getMockTokenRegistry, initMockGetCode, mergeMockSmartContract } from "../testsHelper";
 import { acceptSurrendered } from "./acceptSurrendered";
 
 jest.mock("@govtechsg/token-registry/contracts");
 
 const acceptSurrenderedDocumentParams: TitleEscrowSurrenderDocumentCommand = {
-  tokenRegistry: "0x1122",
-  tokenId: "0x12345",
+  tokenRegistry: "0x0000000000000000000000000000000000000001",
+  tokenId: "0x0000000000000000000000000000000000000000000000000000000000000001",
   network: "goerli",
   dryRun: false,
 };
@@ -21,6 +22,13 @@ describe("title-escrow", () => {
     const mockedConnectERC721: jest.Mock = mockedTradeTrustTokenFactory.connect;
     const mockBurnToken = jest.fn();
     const mockCallStaticBurnToken = jest.fn().mockResolvedValue(undefined);
+    const tokenRegistryAddress = acceptSurrenderedDocumentParams.tokenRegistry;
+    const mockBaseTokenRegistry = getMockTokenRegistry({
+      ownerOfValue: tokenRegistryAddress,
+      address: tokenRegistryAddress,
+      titleEscrowFactoryAddress: AddressZero,
+    });
+    let mockTokenRegistry = mockBaseTokenRegistry;
 
     beforeEach(() => {
       delete process.env.OA_PRIVATE_KEY;
@@ -31,13 +39,15 @@ describe("title-escrow", () => {
         hash: "hash",
         wait: () => Promise.resolve({ transactionHash: "transactionHash" }),
       });
-
-      mockedConnectERC721.mockReturnValue({
+      const mockCustomTokenRegistry = {
         burn: mockBurnToken,
         callStatic: {
           burn: mockCallStaticBurnToken,
         },
-      });
+      };
+      initMockGetCode();
+      mockTokenRegistry = mergeMockSmartContract({ base: mockBaseTokenRegistry, override: mockCustomTokenRegistry });
+      mockedConnectERC721.mockReturnValue(mockTokenRegistry);
       mockBurnToken.mockClear();
       mockCallStaticBurnToken.mockClear();
     });

diff --git a/src/implementations/title-escrow/acceptSurrendered.ts b/src/implementations/title-escrow/acceptSurrendered.ts
@@ -5,7 +5,7 @@ import { BaseTitleEscrowCommand as TitleEscrowSurrenderDocumentCommand } from ".
 
 import { dryRunMode } from "../utils/dryRun";
 import { TransactionReceipt } from "@ethersproject/providers";
-import { TradeTrustToken__factory } from "@govtechsg/token-registry/dist/contracts";
+import { connectToTokenRegistry } from "./helpers";
 
 const { trace } = getLogger("title-escrow:acceptSurrendered");
 
@@ -17,7 +17,7 @@ export const acceptSurrendered = async ({
   ...rest
 }: TitleEscrowSurrenderDocumentCommand): Promise<TransactionReceipt> => {
   const wallet = await getWalletOrSigner({ network, ...rest });
-  const tokenRegistryInstance = await TradeTrustToken__factory.connect(address, wallet);
+  const tokenRegistryInstance = await connectToTokenRegistry({ address, wallet });
   if (dryRun) {
     await dryRunMode({
       estimatedGas: await tokenRegistryInstance.estimateGas.burn(tokenId),

diff --git a/src/implementations/title-escrow/helpers.ts b/src/implementations/title-escrow/helpers.ts
@@ -11,8 +11,23 @@ import { ConnectedSigner } from "../utils/wallet";
 interface ConnectToTitleEscrowArgs {
   tokenId: string;
   address: string;
-  wallet: Wallet | ConnectedSigner;
+  wallet: UserWallet;
 }
+interface ConnectToTokenRegistryArgs {
+  address: string;
+  wallet: UserWallet;
+}
+
+type UserWallet = Wallet | ConnectedSigner;
+
+export const assertAddressIsSmartContract = async (
+  address: string,
+  account: Wallet | ConnectedSigner
+): Promise<void> => {
+  const code = await account.provider.getCode(address);
+  const isContract = code !== "0x" && code !== "0x0"; // Ganache uses 0x0 instead
+  if (!isContract) throw new Error(`Address ${address} is not a valid Contract`);
+};
 
 export const connectToTitleEscrow = async ({
   tokenId,
@@ -24,6 +39,18 @@ export const connectToTitleEscrow = async ({
   return await TitleEscrow__factory.connect(titleEscrowAddress, wallet);
 };
 
+export const connectToTokenRegistry = async ({
+  address,
+  wallet,
+}: ConnectToTokenRegistryArgs): Promise<TradeTrustToken> => {
+  await assertAddressIsSmartContract(address, wallet);
+  const tokenRegistryInstance: TradeTrustToken = await TradeTrustToken__factory.connect(address, wallet);
+  // const isTokenRegistry = await supportsInterface(tokenRegistryInstance, "0x8a198f04")
+  // if(!isTokenRegistry) throw new Error(`Address ${address} is not a supported token registry contract`)
+  await tokenRegistryInstance.callStatic.genesis();
+  return tokenRegistryInstance;
+};
+
 interface validateEndorseChangeOwnerArgs {
   newHolder: string;
   newOwner: string;

diff --git a/src/implementations/title-escrow/rejectSurrendered.test.ts b/src/implementations/title-escrow/rejectSurrendered.test.ts
@@ -1,14 +1,14 @@
-import { TitleEscrow__factory, TradeTrustToken__factory } from "@govtechsg/token-registry/contracts";
+import { TradeTrustToken__factory } from "@govtechsg/token-registry/contracts";
 import { Wallet } from "ethers";
-
 import { BaseTitleEscrowCommand as TitleEscrowSurrenderDocumentCommand } from "../../commands/title-escrow/title-escrow-command.type";
 import { rejectSurrendered } from "./rejectSurrendered";
+import { getMockTokenRegistry, initMockGetCode, mergeMockSmartContract } from "../testsHelper";
 
 jest.mock("@govtechsg/token-registry/contracts");
 
 const rejectSurrenderedDocumentParams: TitleEscrowSurrenderDocumentCommand = {
-  tokenRegistry: "0x1122",
-  tokenId: "0x12345",
+  tokenRegistry: "0x0000000000000000000000000000000000000001",
+  tokenId: "0x0000000000000000000000000000000000000000000000000000000000000001",
   network: "goerli",
   dryRun: false,
 };
@@ -19,63 +19,37 @@ describe("title-escrow", () => {
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore mock static method
     const mockedConnectERC721: jest.Mock = mockedTradeTrustTokenFactory.connect;
-    const mockedTitleEscrowFactory: jest.Mock<TitleEscrow__factory> = TitleEscrow__factory as any;
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore mock static method
-    const mockedConnectTitleEscrowFactory: jest.Mock = mockedTitleEscrowFactory.connect;
-
-    const mockedBeneficiary = jest.fn();
-    const mockedHolder = jest.fn();
     const mockRestoreTitle = jest.fn();
-    const mockTransferEvent = jest.fn();
-    const mockQueryFilter = jest.fn();
     const mockCallStaticRestoreTitle = jest.fn().mockResolvedValue(undefined);
 
-    const mockedLastTitleEscrowAddress = "0xMockedLastTitleEscrowAddress";
-    const mockedLastBeneficiary = "0xMockedLastBeneficiaryAddress";
-    const mockedLastHolder = "0xMockedLastHolderAddress";
+    initMockGetCode();
+
+    const mockBaseTokenRegistry = getMockTokenRegistry({
+      ownerOfValue: rejectSurrenderedDocumentParams.tokenRegistry,
+      address: rejectSurrenderedDocumentParams.tokenRegistry,
+    });
+
+    const mockCustomTokenRegistry = {
+      restore: mockRestoreTitle,
+      callStatic: {
+        restore: mockCallStaticRestoreTitle,
+      },
+    };
+    const mockTokenRegistry = mergeMockSmartContract({
+      base: mockBaseTokenRegistry,
+      override: mockCustomTokenRegistry,
+    });
+
+    mockRestoreTitle.mockReturnValue({
+      hash: "hash",
+      wait: () => Promise.resolve({ transactionHash: "transactionHash" }),
+    });
 
     beforeEach(() => {
       delete process.env.OA_PRIVATE_KEY;
       mockedTradeTrustTokenFactory.mockReset();
       mockedConnectERC721.mockReset();
-      mockedTitleEscrowFactory.mockReset();
-      mockedConnectTitleEscrowFactory.mockReset();
-
-      mockedBeneficiary.mockReturnValue(mockedLastBeneficiary);
-      mockedHolder.mockReturnValue(mockedLastHolder);
-      mockRestoreTitle.mockReturnValue({
-        hash: "hash",
-        wait: () => Promise.resolve({ transactionHash: "transactionHash" }),
-      });
-      mockTransferEvent.mockReturnValue({
-        address: "0x1122",
-        topics: ["0x00000", null, null, "0x12345"],
-      });
-      mockQueryFilter.mockReturnValue([
-        {
-          args: [mockedLastTitleEscrowAddress, "0x1122"],
-        },
-      ]);
-
-      mockedConnectTitleEscrowFactory.mockReturnValue({
-        beneficiary: mockedBeneficiary,
-        holder: mockedHolder,
-      });
-      mockedConnectERC721.mockReturnValue({
-        restore: mockRestoreTitle,
-        filters: { Transfer: mockTransferEvent },
-        queryFilter: mockQueryFilter,
-        callStatic: {
-          restore: mockCallStaticRestoreTitle,
-        },
-      });
-      mockedBeneficiary.mockClear();
-      mockedHolder.mockClear();
-      mockRestoreTitle.mockClear();
-      mockTransferEvent.mockClear();
-      mockQueryFilter.mockClear();
-      mockCallStaticRestoreTitle.mockClear();
+      mockedConnectERC721.mockReturnValue(mockTokenRegistry);
     });
 
     it("should pass in the correct params and successfully rejects a surrendered transferable record", async () => {

diff --git a/src/implementations/title-escrow/rejectSurrendered.ts b/src/implementations/title-escrow/rejectSurrendered.ts
@@ -1,10 +1,10 @@
-import { TradeTrustToken, TradeTrustToken__factory } from "@govtechsg/token-registry/contracts";
 import signale from "signale";
 import { getLogger } from "../../logger";
 import { getWalletOrSigner } from "../utils/wallet";
 import { BaseTitleEscrowCommand as TitleEscrowSurrenderDocumentCommand } from "../../commands/title-escrow/title-escrow-command.type";
 import { dryRunMode } from "../utils/dryRun";
 import { TransactionReceipt } from "@ethersproject/providers";
+import { connectToTokenRegistry } from "./helpers";
 
 const { trace } = getLogger("title-escrow:acceptSurrendered");
 
@@ -16,7 +16,7 @@ export const rejectSurrendered = async ({
   ...rest
 }: TitleEscrowSurrenderDocumentCommand): Promise<TransactionReceipt> => {
   const wallet = await getWalletOrSigner({ network, ...rest });
-  const tokenRegistryInstance: TradeTrustToken = await TradeTrustToken__factory.connect(address, wallet);
+  const tokenRegistryInstance = await connectToTokenRegistry({ address, wallet });
   if (dryRun) {
     await dryRunMode({
       estimatedGas: await tokenRegistryInstance.estimateGas.restore(tokenId),