From d84af922b493de57ab3456fa22d8ef03568cfaf4 Mon Sep 17 00:00:00 2001
From: 0xGorilla <0xGorilla@protonmail.com>
Date: Tue, 7 Sep 2021 16:50:08 +0200
Subject: [PATCH] feat: mocks can be deployed with libs

---
 src/factories/smock-contract.ts       |  6 +++--
 src/index.ts                          | 10 ++++---
 src/sandbox.ts                        | 10 ++++---
 test/contracts/mock/Librarian.sol     | 10 +++++++
 test/contracts/mock/TestLibrary.sol   |  8 ++++++
 test/unit/mock/initialization.spec.ts | 39 +++++++++++++++++++++++++++
 6 files changed, 75 insertions(+), 8 deletions(-)
 create mode 100644 test/contracts/mock/Librarian.sol
 create mode 100644 test/contracts/mock/TestLibrary.sol
 create mode 100644 test/unit/mock/initialization.spec.ts

diff --git a/src/factories/smock-contract.ts b/src/factories/smock-contract.ts
index 651a1c9..54d6c09 100644
--- a/src/factories/smock-contract.ts
+++ b/src/factories/smock-contract.ts
@@ -1,4 +1,5 @@
 import Message from '@nomiclabs/ethereumjs-vm/dist/evm/message';
+import { FactoryOptions } from '@nomiclabs/hardhat-ethers/types';
 import { BaseContract, ContractFactory, ethers } from 'ethers';
 import { Interface } from 'ethers/lib/utils';
 import { ethers as hardhatEthers } from 'hardhat';
@@ -33,9 +34,10 @@ export async function createFakeContract<Contract extends BaseContract>(
 
 export async function createMockContractFactory<T extends ContractFactory>(
   vm: ObservableVM,
-  contractName: string
+  contractName: string,
+  signerOrOptions?: ethers.Signer | FactoryOptions
 ): Promise<MockContractFactory<T>> {
-  const factory = (await hardhatEthers.getContractFactory(contractName)) as unknown as MockContractFactory<T>;
+  const factory = (await hardhatEthers.getContractFactory(contractName, signerOrOptions)) as unknown as MockContractFactory<T>;
 
   const realDeploy = factory.deploy;
   factory.deploy = async (...args: Parameters<T['deploy']>) => {
diff --git a/src/index.ts b/src/index.ts
index 74a6037..9bb7ce4 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,4 +1,5 @@
-import { BaseContract, ContractFactory } from 'ethers';
+import { FactoryOptions } from '@nomiclabs/hardhat-ethers/types';
+import { BaseContract, ContractFactory, ethers } from 'ethers';
 import hre from 'hardhat';
 import { matchers } from './chai-plugin/matchers';
 import { Sandbox } from './sandbox';
@@ -12,10 +13,13 @@ async function fake<T extends BaseContract>(spec: FakeContractSpec, opts: FakeCo
   return await sandbox.fake(spec, opts);
 }
 
-async function mock<T extends ContractFactory>(contractName: string): Promise<MockContractFactory<T>> {
+async function mock<T extends ContractFactory>(
+  contractName: string,
+  signerOrOptions?: ethers.Signer | FactoryOptions
+): Promise<MockContractFactory<T>> {
   if (!sandbox) await init();
 
-  return await sandbox.mock(contractName);
+  return await sandbox.mock(contractName, signerOrOptions);
 }
 
 async function init() {
diff --git a/src/sandbox.ts b/src/sandbox.ts
index d96f43a..6f1a0e9 100644
--- a/src/sandbox.ts
+++ b/src/sandbox.ts
@@ -1,5 +1,6 @@
 import VM from '@nomiclabs/ethereumjs-vm';
-import { BaseContract, ContractFactory } from 'ethers';
+import { FactoryOptions } from '@nomiclabs/hardhat-ethers/types';
+import { BaseContract, ContractFactory, ethers } from 'ethers';
 import hre from 'hardhat';
 import { ethersInterfaceFromSpec } from './factories/ethers-interface';
 import { createFakeContract, createMockContractFactory } from './factories/smock-contract';
@@ -44,8 +45,11 @@ export class Sandbox {
     );
   }
 
-  async mock<T extends ContractFactory>(contractName: string): Promise<MockContractFactory<T>> {
-    return createMockContractFactory(this.vm, contractName);
+  async mock<T extends ContractFactory>(
+    contractName: string,
+    signerOrOptions?: ethers.Signer | FactoryOptions
+  ): Promise<MockContractFactory<T>> {
+    return createMockContractFactory(this.vm, contractName, signerOrOptions);
   }
 
   static async create(): Promise<Sandbox> {
diff --git a/test/contracts/mock/Librarian.sol b/test/contracts/mock/Librarian.sol
new file mode 100644
index 0000000..e9f3fda
--- /dev/null
+++ b/test/contracts/mock/Librarian.sol
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import './TestLibrary.sol';
+
+contract Librarian {
+  function getLibValue() external pure returns (uint256) {
+    return TestLibrary.getSomeValue();
+  }
+}
diff --git a/test/contracts/mock/TestLibrary.sol b/test/contracts/mock/TestLibrary.sol
new file mode 100644
index 0000000..7052cd4
--- /dev/null
+++ b/test/contracts/mock/TestLibrary.sol
@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+library TestLibrary {
+  function getSomeValue() external pure returns (uint256) {
+    return 10;
+  }
+}
diff --git a/test/unit/mock/initialization.spec.ts b/test/unit/mock/initialization.spec.ts
new file mode 100644
index 0000000..d600942
--- /dev/null
+++ b/test/unit/mock/initialization.spec.ts
@@ -0,0 +1,39 @@
+import { BigNumber } from '@ethersproject/bignumber';
+import { smock } from '@src';
+import { Librarian__factory, TestLibrary__factory } from '@typechained';
+import chai, { expect } from 'chai';
+import { ethers } from 'hardhat';
+
+chai.use(smock.matchers);
+
+describe('Mock: Initialization', () => {
+  it('should be able to use libraries', async () => {
+    const testLibrary = await (await ethers.getContractFactory('TestLibrary')).deploy();
+    const librarian = await (
+      await smock.mock<Librarian__factory>('Librarian', {
+        libraries: {
+          TestLibrary: testLibrary.address,
+        },
+      })
+    ).deploy();
+
+    expect(await librarian.getLibValue()).to.equal(10);
+  });
+
+  // TODO: make it work
+  it.skip('should be able to use mocked libraries', async () => {
+    const testLibrary = await (await smock.mock<TestLibrary__factory>('TestLibrary')).deploy();
+    const librarian = await (
+      await smock.mock<Librarian__factory>('Librarian', {
+        libraries: {
+          TestLibrary: testLibrary.address,
+        },
+      })
+    ).deploy();
+
+    const mockValue = BigNumber.from(123);
+    testLibrary.getSomeValue.returns(mockValue);
+
+    expect(await librarian.getLibValue()).to.equal(mockValue);
+  });
+});