diff --git a/assembly/__tests__/vm-mock.spec.ts b/assembly/__tests__/vm-mock.spec.ts index a0745d0a..2220507a 100644 --- 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 index 523b4741..2f6cfc33 100644 --- 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), }); },