feat: 🎸 getTokensByPartition with test

Author: Victor Wiebe

src/contract_wrappers/modules/transfer_manager/__tests__/lock_up_transfer_manager_wrapper.test.ts

@@ -20,6 +20,7 @@ import {
   valueToWei,
   weiToValue,
 } from '../../../../utils/convert';
+import { Partition } from '../../../../types';
 
 describe('LockUpTransferManagerWrapper', () => {
   let target: LockUpTransferManagerWrapper;
@@ -511,6 +512,64 @@ describe('LockUpTransferManagerWrapper', () => {
     });
   });
 
+  describe('getTokensByPartition', () => {
+    test('should call to getTokensByPartition', async () => {
+      const expectedDecimalsResult = new BigNumber(18);
+      const expectedResult = valueToWei(new BigNumber(100), expectedDecimalsResult);
+      const mockedParams = {
+        partition: Partition.Unlocked,
+        tokenHolder: '0x8888888888888888888888888888888888888888',
+        additionalBalance: new BigNumber(10),
+      };
+
+      // Security Token Address expected
+      const expectedSecurityTokenAddress = '0x3333333333333333333333333333333333333333';
+      // Setup get Security Token Address
+      const mockedGetSecurityTokenAddressMethod = mock(MockedCallMethod);
+      when(mockedContract.securityToken).thenReturn(instance(mockedGetSecurityTokenAddressMethod));
+      when(mockedGetSecurityTokenAddressMethod.callAsync()).thenResolve(expectedSecurityTokenAddress);
+      when(mockedContractFactory.getSecurityTokenContract(expectedSecurityTokenAddress)).thenResolve(
+        instance(mockedSecurityTokenContract),
+      );
+      const mockedSecurityTokenDecimalsMethod = mock(MockedCallMethod);
+      when(mockedSecurityTokenDecimalsMethod.callAsync()).thenResolve(expectedDecimalsResult);
+      when(mockedSecurityTokenContract.decimals).thenReturn(instance(mockedSecurityTokenDecimalsMethod));
+
+      // Mocked method
+      const mockedMethod = mock(MockedCallMethod);
+      // Stub the method
+      when(mockedContract.getTokensByPartition).thenReturn(instance(mockedMethod));
+      // Stub the request
+      when(
+        mockedMethod.callAsync(
+          mockedParams.partition,
+          mockedParams.tokenHolder,
+          objectContaining(valueToWei(mockedParams.additionalBalance, expectedDecimalsResult)),
+        ),
+      ).thenResolve(expectedResult);
+
+      // Real call
+      const result = await target.getTokensByPartition(mockedParams);
+      // Result expectation
+      expect(result).toEqual(weiToValue(expectedResult, expectedDecimalsResult));
+
+      // Verifications
+      verify(mockedContract.getTokensByPartition).once();
+      verify(
+        mockedMethod.callAsync(
+          mockedParams.partition,
+          mockedParams.tokenHolder,
+          objectContaining(valueToWei(mockedParams.additionalBalance, expectedDecimalsResult)),
+        ),
+      ).once();
+      verify(mockedContract.securityToken).once();
+      verify(mockedGetSecurityTokenAddressMethod.callAsync()).once();
+      verify(mockedContractFactory.getSecurityTokenContract(expectedSecurityTokenAddress)).once();
+      verify(mockedSecurityTokenDecimalsMethod.callAsync()).once();
+      verify(mockedSecurityTokenContract.decimals).once();
+    });
+  });
+
   describe('verifyTransfer', () => {
     test('should verify Transfer', async () => {
       const statusCode = new BigNumber(2);

src/contract_wrappers/modules/transfer_manager/lock_up_transfer_manager_wrapper.ts

@@ -23,6 +23,7 @@ import {
   EventCallback,
   GetLogs,
   GetLogsAsyncParams,
+  Partition,
   Subscribe,
   SubscribeAsyncParams,
   TransferResult,
@@ -139,6 +140,12 @@ interface UserAddressParams {
   user: string;
 }
 
+interface GetTokensByPartitionParams {
+  partition: Partition;
+  tokenHolder: string;
+  additionalBalance: BigNumber;
+}
+
 interface VerifyTransferParams {
   from: string;
   to: string;
@@ -289,6 +296,22 @@ export default class LockUpTransferManagerWrapper extends ModuleWrapper {
     return weiToValue(await (await this.contract).getLockedTokenToUser.callAsync(params.user), decimals);
   };
 
+  /**
+   * getTokensByPartition
+   */
+  public getTokensByPartition = async (params: GetTokensByPartitionParams): Promise<BigNumber> => {
+    assert.isNonZeroETHAddressHex('Token Holder', params.tokenHolder);
+    const decimals = await (await this.securityTokenContract()).decimals.callAsync();
+    return weiToValue(
+      await (await this.contract).getTokensByPartition.callAsync(
+        params.partition,
+        params.tokenHolder,
+        valueToWei(params.additionalBalance, decimals),
+      ),
+      decimals,
+    );
+  };
+
   /*
    * verifyTransfer
    */