From 958395df7a1558c09a9f169e805b2847312f13fc Mon Sep 17 00:00:00 2001
From: Mateusz Marszalek <matti.marszalek@gmail.com>
Date: Mon, 25 Nov 2024 21:17:10 +0100
Subject: [PATCH] Test refactor to new async/await

---
 test/func_blocks/fu/test_lza.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/test/func_blocks/fu/test_lza.py b/test/func_blocks/fu/test_lza.py
index 6682f8231..9c5f0cb97 100644
--- a/test/func_blocks/fu/test_lza.py
+++ b/test/func_blocks/fu/test_lza.py
@@ -1,6 +1,6 @@
+import random
 from coreblocks.func_blocks.fu.fpu.lza import *
 from coreblocks.func_blocks.fu.fpu.fpu_common import FPUParams
-from random import randint
 from transactron import TModule
 from transactron.lib import AdapterTrans
 from transactron.testing import *
@@ -50,19 +50,19 @@ def test_manual(self):
         help_values = TestLZA.HelpValues(params)
         lza = TestLZA.LZAModuleTest(params)
 
-        def random_test(seed, iters):
+        async def random_test(sim: TestbenchContext, seed: int, iters: int):
             xor_mask = (2**params.sig_width) - 1
             random.seed(seed)
             for _ in range(iters):
-                sig_a = randint(1 << (params.sig_width - 1), (2**params.sig_width) - 1)
-                sig_b = randint(1 << (params.sig_width - 1), sig_a)
+                sig_a = random.randint(1 << (params.sig_width - 1), (2**params.sig_width) - 1)
+                sig_b = random.randint(1 << (params.sig_width - 1), sig_a)
                 sig_b = (sig_b ^ xor_mask) | (1 << params.sig_width)
-                resp = yield from lza.predict_request_adapter.call({"sig_a": sig_a, "sig_b": sig_b, "carry": 0})
+                resp = await lza.predict_request_adapter.call(sim, {"sig_a": sig_a, "sig_b": sig_b, "carry": 0})
                 pred_lz = resp["shift_amount"]
                 true_lz = clz(sig_a, sig_b, 0, params.sig_width)
                 assert pred_lz == true_lz or (pred_lz + 1) == true_lz
 
-        def lza_test():
+        async def lza_test(sim: TestbenchContext):
             test_cases = [
                 {
                     "sig_a": help_values.test_val_sig_a_1,
@@ -140,13 +140,13 @@ def lza_test():
                 {"shift_amount": 7, "is_zero": 0},
             ]
             for i in range(len(test_cases)):
-                resp = yield from lza.predict_request_adapter.call(test_cases[i])
+                resp = await lza.predict_request_adapter.call(sim, test_cases[i])
                 assert resp["shift_amount"] == expected_results[i]["shift_amount"]
                 assert resp["is_zero"] == expected_results[i]["is_zero"]
 
-        def test_process():
-            yield from lza_test()
-            yield from random_test(2024, 20)
+        async def test_process(sim: TestbenchContext):
+            await lza_test(sim)
+            await random_test(sim, 2024, 20)
 
         with self.run_simulation(lza) as sim:
-            sim.add_process(test_process)
+            sim.add_testbench(test_process)