vm-mock: add current slot mock

massalabs · Sep 19, 2024 · 69bb9f4 · 69bb9f4
1 parent 01baac3
commit 69bb9f4
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/assembly/__tests__/env-helpers.spec.ts b/assembly/__tests__/env-helpers.spec.ts
@@ -2,7 +2,7 @@
 
 import { env } from '../env';
 import { Address } from '../std';
-import { addAddressToLedger } from '../vm-mock';
+import { addAddressToLedger, mockCurrentSlot } from '../vm-mock';
 
 const testAddress = new Address(
  'AU12E6N5BFAdC2wyiBV6VJjqkWhpz1kLVp2XpbRdSnL1mKjCWT6oR',
@@ -17,6 +17,12 @@ describe('Testing env coins related functions', () => {
  expect(env.currentThread()).toBe(1);
  });
 
+ test('current slot', () => {
+ mockCurrentSlot(1234, 8);
+ expect(env.currentPeriod()).toBe(1234);
+ expect(env.currentThread()).toBe(8);
+ });
+
  test('setBytecode', () => {
  env.setBytecode(new StaticArray(0));
  const byteCode = env.getBytecode().toString();

diff --git a/assembly/vm-mock/env.ts b/assembly/vm-mock/env.ts
@@ -182,3 +182,6 @@ export declare function mockSetChainId(value: number): void;
  */
 @external("massa", "assembly_script_set_origin_operation_id")
 export declare function mockOriginOperationId(opId: string): void;
+
+@external("massa", "assembly_script_set_slot")
+export declare function mockCurrentSlot(period: u64, thread: u8): void;
diff --git a/vm-mock/vm.js b/vm-mock/vm.js
@@ -84,6 +84,8 @@ let adminContext = false;
 
 let webModule;
 
+let currentSlot = { period: 0n, thread: 1 };
+
 const scCallMockStack = [];
 let callCoins = 0n; // Default value, coins for a call
 let spentCoins = 0n; // Coins spent during the call
@@ -698,11 +700,18 @@ export default function createMockedABI(
  },
 
  assembly_script_get_current_period() {
- return BigInt(0);
+ return currentSlot.period;
  },
 
  assembly_script_get_current_thread() {
- return 1;
+ return currentSlot.thread;
+ },
+
+ assembly_script_set_slot(period, thread) {
+ currentSlot = {
+ period: BigInt(period),
+ thread,
+ };
  },
 
  assembly_script_set_bytecode(bytecodePtr) {