From 8c46ca5bffe675f3e0c7edb849d630bf2410bbb8 Mon Sep 17 00:00:00 2001
From: Pierre Seznec <petterseznec@gmail.com>
Date: Thu, 19 Sep 2024 13:37:43 +0200
Subject: [PATCH] vm-mock: add current slot mock

---
 assembly/__tests__/env-helpers.spec.ts |  8 +++++++-
 assembly/vm-mock/env.ts                |  3 +++
 vm-mock/vm.js                          | 13 +++++++++++--
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/assembly/__tests__/env-helpers.spec.ts b/assembly/__tests__/env-helpers.spec.ts
index e7952547..e3fbb249 100644
--- 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
index 2b60bac2..dfbb189d 100644
--- 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
index cc3cd4d3..e82a2205 100644
--- 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) {