vm-mock: fix transfer mock sender

massalabs · Feb 29, 2024 · c8d9a42 · c8d9a42
1 parent ac3c14e
commit c8d9a42
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 21 deletions.
diff --git a/assembly/__tests__/vm-mock.spec.ts b/assembly/__tests__/vm-mock.spec.ts
@@ -13,6 +13,7 @@ import {
   getOriginOperationId,
   balance,
   balanceOf,
+  transferCoins,
 } from '../std';
 import { changeCallStack, resetStorage } from '../vm-mock/storage';
 import {
@@ -39,6 +40,14 @@ const testAddress2 = new Address(
   'AS12E6N5BFAdC2wyiBV6VJjqkWhpz1kLVp2XpbRdSnL1mKjCWT6oP',
 );
 
+const callerAddress: Address = new Address(
+  'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKq',
+);
+
+const contractAddress: Address = new Address(
+  'AS12BqZEQ6sByhRLyEuf0YbQmcF2PsDdkNNG1akBJu9XcjZA1eT',
+);
+
 const keyTest = 'test';
 const valueTest = 'value';
 
@@ -170,13 +179,6 @@ describe('Testing mocked Storage and CallStack', () => {
   });
 });
 
-const callerAddress: Address = new Address(
-  'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKq',
-);
-const contractAddress: Address = new Address(
-  'AS12BqZEQ6sByhRLyEuf0YbQmcF2PsDdkNNG1akBJu9XcjZA1eT',
-);
-
 describe('Testing mocked Context', () => {
   beforeEach(() => {
     mockAdminContext(false);
@@ -347,3 +349,37 @@ describe('balance mock', () => {
     expect(Storage.get('thekey')).toBe('thevalue');
   });
 });
+
+describe('transfer coins', () => {
+  const amount = 666;
+  it('contract transfer amount to user', () => {
+    const initialContractBalance = balance();
+
+    transferCoins(testAddress, amount);
+    expect(balanceOf(testAddress.toString())).toBe(amount);
+    expect(balance()).toBe(initialContractBalance - amount);
+    expect(balanceOf(contractAddress.toString())).toBe(
+      initialContractBalance - amount,
+    );
+  });
+
+  it('transfer increase user balance', () => {
+    const initialUserBalance = 123;
+    const initialContractBalance = balance();
+
+    mockBalance(testAddress.toString(), initialUserBalance);
+
+    transferCoins(testAddress, amount);
+    expect(balanceOf(testAddress.toString())).toBe(amount + initialUserBalance);
+    expect(balance()).toBe(initialContractBalance - amount);
+    expect(balanceOf(contractAddress.toString())).toBe(
+      initialContractBalance - amount,
+    );
+  });
+
+  throws('throw if insufficient balance', () => {
+    mockBalance(contractAddress.toString(), 0);
+    expect(balance()).toBe(0);
+    transferCoins(testAddress, amount);
+  });
+});
diff --git a/vm-mock/vm.js b/vm-mock/vm.js
@@ -24,7 +24,6 @@ function ERROR(msg) {
   throw new Error(msg);
 }
 
-
 /**
  * return a random string
  *
@@ -470,10 +469,6 @@ export default function createMockedABI(
         adminContext = isAdmin;
       },
 
-      assembly_script_mock_not_admin_context() {
-        adminContext = false;
-      },
-
       assembly_script_set_deploy_context(addrPtr) {
         adminContext = true;
         // Ensure the caller address is different from the contract address
@@ -682,22 +677,31 @@ export default function createMockedABI(
         return newArrayBuffer(addressLedger.contract);
       },
 
-      assembly_script_transfer_coins(_addressPtr, _coinsAmount) {
-        const address = ptrToString(_addressPtr);
-        if (!ledger.has(address)) {
-          ledger.set(address, {
+      assembly_script_transfer_coins(addressPtr, amount) {
+        const toAddress = ptrToString(addressPtr);
+
+        if (!ledger.has(toAddress)) {
+          ledger.set(toAddress, {
             storage: new Map(),
             contract: '',
             balance: BigInt(0),
           });
         }
-        const callerBalance = ledger.get(callerAddress).balance;
+        // get last elt of callStack
+        const callStackArray = callStack.split(' , ');
+        const fromAddress = callStackArray[callStackArray.length - 1];
+
+        if (!ledger.has(fromAddress)) {
+          ERROR(`Sending address ${fromAddress} does not exist in ledger.`);
+        }
+        const senderBalance = ledger.get(fromAddress).balance;
 
-        if (callerBalance < BigInt(_coinsAmount)) {
-          ERROR('not enough balance to transfer ' + _coinsAmount + ' coins.');
+        if (senderBalance < BigInt(amount)) {
+          ERROR('not enough balance to transfer ' + amount + ' coins.');
         }
-        ledger.get(callerAddress).balance -= BigInt(_coinsAmount);
-        ledger.get(address).balance += BigInt(_coinsAmount);
+
+        ledger.get(toAddress).balance += BigInt(amount);
+        ledger.get(fromAddress).balance -= BigInt(amount);
       },
 
       assembly_script_transfer_coins_for(