feat: mocks can now edit internal variables (defi-wonderland#7)

* fix: mock created from mock factory
* feat: editable storage
* docs: updated docs

docs/source/mocks.rst

@@ -7,7 +7,7 @@ What are they
 Mocks are deployed contract wrappers that have all of the fake's functionality and even more.
 
 Because they are actually deployed contract, they can have actual **logic inside** that can be called through.
-And because they have an state, **internal variable values can be overwritten** 🥳
+And because they have a storage, **internal variable values can be overwritten** 🥳
 
 
 How to use
@@ -37,17 +37,15 @@ How to use
 
   .. code-block:: javascript
 
-    import { MyOracle, Counter, Counter__factory } from '@typechained';
-    import { MockContract, lopt } from '@defi-wonderland/lopt';
-
-    chai.use(lopt.matchers);
+    import { Counter } from '@typechained';
+    import { MockContract, MockContractFactory, lopt } from '@defi-wonderland/lopt';
 
     describe('Counter', () => {
-      let counterFactory: Counter__factory;
+      let counterFactory: MockContractFactory<Counter>;
       let counter: MockContract<Counter>;
 
       before(async () => {
-        counterFactory = (await ethers.getContractFactory('Counter')) as Counter__factory;
+        counterFactory = await lopt.mock<Counter>('Counter');
       });
 
       beforeEach(async () => {
@@ -86,10 +84,19 @@ Internal variables override
 
 .. container:: code-explanation
 
-  **Not yet developed**, but it should look something like this (open to new ideas)
+  Set the value of an internal variable
+
+  .. code-block:: javascript
+
+    await mock.setVariable('_myInternalVariable', true);
+
+.. container:: code-explanation
+
+  Set the value of an internal struct
 
   .. code-block:: javascript
 
-    await mock.shouldGetCrazy(); // returns false
-    await mock._myInternalVariable.set(true);
-    await mock.shouldGetCrazy(); // returns true
+    await mock.setVariable('_myInternalStruct', {
+      _valueA: true,
+      _valueB: 123
+    });

src/factories/lopt-contract.ts

@@ -1,13 +1,16 @@
 import Message from '@nomiclabs/ethereumjs-vm/dist/evm/message';
 import { BaseContract, ethers } from 'ethers';
 import { Interface } from 'ethers/lib/utils';
+import { ethers as hardhatEthers } from 'hardhat';
 import { Observable } from 'rxjs';
 import { distinct, filter, map, share, withLatestFrom } from 'rxjs/operators';
+import { EditableStorageLogic as EditableStorage } from '../logic/editable-storage-logic';
 import { ProgrammableFunctionLogic, SafeProgrammableContract } from '../logic/programmable-function-logic';
 import { ObservableVM } from '../observable-vm';
 import { Sandbox } from '../sandbox';
-import { ContractCall, FakeContract, MockContract, ProgrammableContractFunction, ProgrammedReturnValue } from '../types';
+import { ContractCall, FakeContract, MockContract, MockContractFactory, ProgrammableContractFunction, ProgrammedReturnValue } from '../types';
 import { fromFancyAddress, impersonate, toFancyAddress, toHexString } from '../utils';
+import { getStorageLayout } from '../utils/storage';
 
 export async function createFakeContract<Contract extends BaseContract>(
   vm: ObservableVM,
@@ -16,32 +19,44 @@ export async function createFakeContract<Contract extends BaseContract>(
   provider: ethers.providers.Provider
 ): Promise<FakeContract<Contract>> {
   const fake = await initContract<FakeContract<Contract>>(vm, address, contractInterface, provider);
-  const uniqueFns = getUniqueFunctionNamesBySighash(contractInterface, Object.keys(fake.functions));
-  const fnsToFill: [string | null, string][] = [...Object.entries(uniqueFns), [null, 'fallback']];
+  const contractFunctions = getContractFunctionsNameAndSighash(contractInterface, Object.keys(fake.functions));
 
-  fnsToFill.forEach(([sighash, name]) => {
+  // attach to every contract function, all the programmable and watchable logic
+  contractFunctions.forEach(([sighash, name]) => {
     const { encoder, calls$, results$ } = getFunctionEventData(vm, contractInterface, fake.address, sighash);
-
     const functionLogic = new SafeProgrammableContract(name, calls$, results$, encoder);
     fillProgrammableContractFunction(fake[name], functionLogic);
   });
 
   return fake;
 }
 
-export async function createMockContract<Contract extends BaseContract>(vm: ObservableVM, contract: Contract): Promise<MockContract<Contract>> {
-  const mock = contract as MockContract<Contract>;
-  const uniqueFns = getUniqueFunctionNamesBySighash(mock.interface, Object.keys(mock.functions));
-  const fnsToFill: [string | null, string][] = [...Object.entries(uniqueFns), [null, 'fallback']];
-
-  fnsToFill.forEach(([sighash, name]) => {
-    const { encoder, calls$, results$ } = getFunctionEventData(vm, mock.interface, mock.address, sighash);
-
-    const functionLogic = new ProgrammableFunctionLogic(name, calls$, results$, encoder);
-    fillProgrammableContractFunction(mock[name], functionLogic);
-  });
+export async function createMockContractFactory<Contract extends BaseContract>(
+  vm: ObservableVM,
+  contractName: string
+): Promise<MockContractFactory<Contract>> {
+  const factory = (await hardhatEthers.getContractFactory(contractName)) as MockContractFactory<Contract>;
+
+  const realDeploy = factory.deploy;
+  factory.deploy = async (...args) => {
+    const mock = (await realDeploy.apply(factory, args)) as MockContract<Contract>;
+    const contractFunctions = getContractFunctionsNameAndSighash(mock.interface, Object.keys(mock.functions));
+
+    // attach to every contract function, all the programmable and watchable logic
+    contractFunctions.forEach(([sighash, name]) => {
+      const { encoder, calls$, results$ } = getFunctionEventData(vm, mock.interface, mock.address, sighash);
+      const functionLogic = new ProgrammableFunctionLogic(name, calls$, results$, encoder);
+      fillProgrammableContractFunction(mock[name], functionLogic);
+    });
+
+    // attach to every internal variable, all the editable logic
+    const editableStorage = new EditableStorage(await getStorageLayout(contractName), vm.getManager(), mock.address);
+    mock.setVariable = editableStorage.setVariable.bind(editableStorage);
+
+    return mock;
+  };
 
-  return mock;
+  return factory;
 }
 
 async function initContract<T extends BaseContract>(
@@ -137,17 +152,19 @@ function fillProgrammableContractFunction(fn: ProgrammableContractFunction, logi
  *
  * @param contractInterface contract interface in order to get the sighash of a name
  * @param names function names to be filtered
- * @returns unique function names and its sighashes
+ * @returns array of sighash and function name
  */
-function getUniqueFunctionNamesBySighash(contractInterface: ethers.utils.Interface, names: string[]): { [sighash: string]: string } {
-  let result: { [sighash: string]: string } = {};
+function getContractFunctionsNameAndSighash(contractInterface: ethers.utils.Interface, names: string[]): [string | null, string][] {
+  let functions: { [sighash: string]: string } = {};
+
   names.forEach((name) => {
     const sighash = contractInterface.getSighash(name);
-    if (!result[sighash] || !name.includes('(')) {
-      result[sighash] = name;
+    if (!functions[sighash] || !name.includes('(')) {
+      functions[sighash] = name;
     }
   });
-  return result;
+
+  return [...Object.entries(functions), [null, 'fallback']];
 }
 
 function parseMessage(message: Message, contractInterface: Interface, sighash: string | null): ContractCall {

src/index.ts

@@ -1,7 +1,7 @@
 import { BaseContract } from 'ethers';
 import { matchers } from './chai-plugin/matchers';
 import { Sandbox } from './sandbox';
-import { FakeContract, FakeContractOptions, FakeContractSpec, MockContract } from './types';
+import { FakeContract, FakeContractOptions, FakeContractSpec, MockContractFactory } from './types';
 
 let sandbox: Sandbox;
 
@@ -10,9 +10,9 @@ async function fake<Type extends BaseContract>(spec: FakeContractSpec, opts: Fak
   return await sandbox.fake(spec, opts);
 }
 
-async function mock<Contract extends BaseContract>(contract: Contract): Promise<MockContract<Contract>> {
+async function mock<Contract extends BaseContract>(contractName: string): Promise<MockContractFactory<Contract>> {
   if (!sandbox) sandbox = await Sandbox.create();
-  return await sandbox.mock(contract);
+  return await sandbox.mock(contractName);
 }
 
 export * from './types';

src/logic/editable-storage-logic.ts

@@ -0,0 +1,28 @@
+import { LoptVMManager } from '../types';
+import { convertToStorageSlots, fromHexString, toFancyAddress } from '../utils';
+
+export class EditableStorageLogic {
+  private storageLayout: any;
+  private contractAddress: string;
+  private vmManager: LoptVMManager;
+
+  constructor(storageLayout: any, vmManager: LoptVMManager, contractAddress: string) {
+    this.storageLayout = storageLayout;
+    this.vmManager = vmManager;
+    this.contractAddress = contractAddress;
+  }
+
+  async setVariable(variableName: string, value: any) {
+    if (value === undefined || value === null) return;
+
+    const slots = convertToStorageSlots(this.storageLayout, variableName, value);
+
+    for (const slot of slots) {
+      await this.vmManager.putContractStorage(
+        toFancyAddress(this.contractAddress),
+        fromHexString(slot.hash.toLowerCase()),
+        fromHexString(slot.value)
+      );
+    }
+  }
+}

src/sandbox.ts

@@ -2,9 +2,9 @@ import VM from '@nomiclabs/ethereumjs-vm';
 import { BaseContract } from 'ethers';
 import hre from 'hardhat';
 import { ethersInterfaceFromSpec } from './factories/ethers-interface';
-import { createFakeContract, createMockContract } from './factories/lopt-contract';
+import { createFakeContract, createMockContractFactory } from './factories/lopt-contract';
 import { ObservableVM } from './observable-vm';
-import { FakeContract, FakeContractOptions, FakeContractSpec, MockContract } from './types';
+import { FakeContract, FakeContractOptions, FakeContractSpec, MockContractFactory } from './types';
 import { getHardhatBaseProvider, makeRandomAddress } from './utils';
 
 // Handle hardhat ^2.4.0
@@ -44,8 +44,8 @@ export class Sandbox {
     );
   }
 
-  async mock<Contract extends BaseContract>(contract: Contract): Promise<MockContract<Contract>> {
-    return createMockContract(this.vm, contract);
+  async mock<Contract extends BaseContract>(contractName: string): Promise<MockContractFactory<Contract>> {
+    return createMockContractFactory(this.vm, contractName);
   }
 
   static async create(): Promise<Sandbox> {

src/types.ts

@@ -4,6 +4,7 @@ import { Provider } from '@ethersproject/abstract-provider';
 import { Signer } from '@ethersproject/abstract-signer';
 import { BaseContract, Contract, ContractFactory, ethers } from 'ethers';
 import { Artifact } from 'hardhat/types';
+import { EditableStorageLogic } from './logic/editable-storage-logic';
 import { WatchableFunctionLogic } from './logic/watchable-function-logic';
 
 export type FakeContractSpec = Artifact | Contract | ContractFactory | ethers.utils.Interface | string | (JsonFragment | Fragment | string)[];
@@ -17,6 +18,8 @@ export type ProgrammedReturnValue = any;
 
 export interface LoptVMManager {
   putContractCode: (address: Buffer, code: Buffer) => Promise<void>;
+  getContractStorage: (address: Buffer, slotHash: Buffer) => Promise<Buffer>;
+  putContractStorage: (address: Buffer, slotHash: Buffer, slotValue: Buffer) => Promise<void>;
 }
 
 export interface WatchableContractFunction {
@@ -55,4 +58,9 @@ export type MockContract<Contract extends BaseContract> = BaseContract &
   } & {
     wallet: Signer;
     fallback: ProgrammableContractFunction;
+    setVariable: EditableStorageLogic['setVariable'];
   };
+
+export interface MockContractFactory<Contract extends BaseContract> extends ContractFactory {
+  deploy: (...args: Array<any>) => Promise<MockContract<Contract>>;
+}

src/utils/hardhat.ts

@@ -57,7 +57,7 @@ export const getHardhatBaseProvider = async (runtime: HardhatRuntimeEnvironment)
  * @param address String address to convert into the fancy new address type.
  * @returns Fancified address.
  */
-export const toFancyAddress = (address: string): any => {
+export const toFancyAddress = (address: string): Buffer => {
   const fancyAddress = fromHexString(address);
   (fancyAddress as any).buf = fromHexString(address);
   (fancyAddress as any).toString = (encoding?: any) => {

src/utils/index.ts

@@ -3,4 +3,5 @@ export * from './hardhat';
 export * from './hex-utils';
 export * from './misc';
 export * from './serdes';
+export * from './storage';
 export * from './wallet';

src/utils/misc.ts

@@ -1,4 +1,33 @@
+import { BigNumber } from 'ethers';
+import { isPojo } from './serdes';
+
 const timeWords = [null, 'once', 'twice', 'thrice'];
 export function humanizeTimes(count: number) {
   return timeWords[count] || `${count || 0} times`;
 }
+
+/**
+ * Custom object flatten
+ *
+ * @param obj Object to flatten.
+ * @param prefix Current object prefix (used recursively).
+ * @param result Current result (used recursively).
+ * @returns Flattened object.
+ */
+export function flatten(obj: any, prefix: string = '', result: any = {}): Object {
+  return Object.entries(obj).reduce((acc, [key, val]) => {
+    const subKey = `${prefix}${key}`;
+
+    if (BigNumber.isBigNumber(val) || typeof val === 'string' || typeof val === 'number' || typeof val === 'boolean') {
+      result[subKey] = val;
+    } else if (Array.isArray(val)) {
+      val.forEach((valItem, index) => flatten(valItem, !prefix ? `${index}` : `${prefix}${index}.`, acc));
+    } else if (isPojo(val)) {
+      flatten(val, `${subKey}.`, acc);
+    } else {
+      throw new Error('Cannot flatten unsupported object type');
+    }
+
+    return acc;
+  }, result);
+}