vm-mock: fix balance mock

massalabs · Feb 29, 2024 · 4ffdaf0 · 4ffdaf0
1 parent d18b384
commit 4ffdaf0
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/assembly/__tests__/vm-mock.spec.ts b/assembly/__tests__/vm-mock.spec.ts
@@ -11,10 +11,13 @@ import {
   isEvmSignatureValid,
   evmGetPubkeyFromSignature,
   getOriginOperationId,
+  balance,
+  balanceOf,
 } from '../std';
 import { changeCallStack, resetStorage } from '../vm-mock/storage';
 import {
   mockAdminContext,
+  mockBalance,
   mockOriginOperationId,
   mockSetChainId,
   setDeployContext,
@@ -325,3 +328,22 @@ describe('Testing mocked origin operation id', () => {
     expect(opId).not.toBe(mockedOpId);
   });
 });
+
+describe('balance mock', () => {
+  it('mock balance', () => {
+    mockBalance(testAddress.toString(), 9999);
+    expect(balanceOf(testAddress.toString())).toBe(9999);
+  });
+
+  it('mock contract balance', () => {
+    mockBalance(contractAddress.toString(), 9999);
+    expect(balance()).toBe(9999);
+  });
+
+  it('mock contract balance and preserve storage', () => {
+    Storage.set('thekey', 'thevalue');
+    mockBalance(contractAddress.toString(), 9999);
+    expect(balance()).toBe(9999);
+    expect(Storage.get('thekey')).toBe('thevalue');
+  });
+});
diff --git a/vm-mock/vm.js b/vm-mock/vm.js
@@ -733,7 +733,17 @@ export default function createMockedABI(
 
       assembly_script_mock_balance(aPtr, amount) {
         const addr = ptrToString(aPtr);
+        if (!ledger.has(addr)) {
+          ledger.set(addr, {
+            storage: new Map(),
+            contract: '',
+            balance: BigInt(amount),
+          });
+          return;
+        }
         ledger.set(addr, {
+          storage: ledger.get(addr).storage,
+          contract: ledger.get(addr).contract,
           balance: BigInt(amount),
         });
       },