fix: enable count transfer manager refactor

Author: Victor Wiebe

src/procedures/__tests__/EnableCountTransferManager.ts

@@ -1,6 +1,7 @@
 import { ImportMock, MockManager } from 'ts-mock-imports';
-import { spy, restore } from 'sinon';
+import sinon, { stub, restore } from 'sinon';
 import * as contractWrappersModule from '@polymathnetwork/contract-wrappers';
+import { ModuleName } from '@polymathnetwork/contract-wrappers';
 import * as contextModule from '../../Context';
 import * as wrappersModule from '../../PolymathBase';
 import * as tokenFactoryModule from '../../testUtils/MockedTokenFactoryModule';
@@ -14,6 +15,7 @@ import {
   ProcedureType,
 } from '../../types';
 
+
 const params: EnableCountTransferManagerProcedureArgs = {
   symbol: 'TEST1',
   maxHolderCount: 3,
@@ -62,21 +64,35 @@ describe('EnableCountTransferManager', () => {
 
   describe('EnableCountTransferManager', () => {
     test('should add a transaction to the queue to enable count transfer manager', async () => {
-      const addTransactionSpy = spy(target, 'addTransaction');
+      const addModuleWithLabelArgsSpy = sinon.spy();
+      const addTransactionStub = stub(target, 'addTransaction');
+
       securityTokenMock.mock('addModuleWithLabel', Promise.resolve('AddModuleWithLabel'));
+      const { addModuleWithLabel } = securityTokenMock.getMockInstance();
+      addTransactionStub.withArgs(addModuleWithLabel).returns(addModuleWithLabelArgsSpy);
 
       // Real call
       await target.prepareTransactions();
 
       // Verifications
+      expect(addModuleWithLabelArgsSpy.getCall(0).args[0]).toEqual({
+        moduleName: ModuleName.CountTransferManager,
+        address: moduleFactoryAddress,
+        archived: false,
+        data: {
+          maxHolderCount: params.maxHolderCount,
+        },
+      });
+      expect(addModuleWithLabelArgsSpy.callCount).toEqual(1);
+
       expect(
-        addTransactionSpy
+        addTransactionStub
           .getCall(0)
           .calledWithExactly(securityTokenMock.getMockInstance().addModuleWithLabel, {
             tag: PolyTransactionTag.EnableCountTransferManager,
           })
       ).toEqual(true);
-      expect(addTransactionSpy.callCount).toEqual(1);
+      expect(addTransactionStub.callCount).toEqual(1);
     });
 
     test('should throw if there is no valid security token supplied', async () => {