diff --git a/contracts/tests/mockOneProto.sol b/contracts/tests/mockOneProto.sol
index 18bccac..6311279 100644
--- a/contracts/tests/mockOneProto.sol
+++ b/contracts/tests/mockOneProto.sol
@@ -15,6 +15,10 @@ contract MockConnectOne is ConnectOne {
   function setUint(uint setId, uint val) override internal {}
 
   function sub(uint x, uint y) internal override pure returns (uint z) {
-      z = 100000;
+      z = 21 * 10 ** 18;
   }
+
+    function getOneProtoAddress() internal override view returns (address payable) {
+        return payable(oneProtoAddr);
+    }
 }
diff --git a/test/oneProto.js b/test/oneProto.js
index 1283c61..83d9fcb 100644
--- a/test/oneProto.js
+++ b/test/oneProto.js
@@ -50,7 +50,7 @@ const {
             5,
             0
         ).encodeABI();
-        await mock.givenMethodReturnUint(getExpectedReturn, String(sellAmt));
+        await mock.givenMethodReturn(getExpectedReturn, web3.eth.abi.encodeParameters(["uint256", "uint256[]"], [20000, [0,0,0,1]]));
       // mocking stake
       let swapWithReferral = await oneProtoContract.methods.swapWithReferral(
             mock.address,