fix(tests_id_generation): test ID persistence

ZettaAI · Jan 4, 2024 · c1fd9af · c1fd9af
1 parent f79375a
commit c1fd9af
Showing 1 changed file with 39 additions and 0 deletions.
diff --git a/tests/unit/mazepa/test_id_generation.py b/tests/unit/mazepa/test_id_generation.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import multiprocessing
 from functools import partial
 from typing import Any, Callable, Mapping
 
@@ -180,3 +181,41 @@ def test_generate_invocation_id_flow_schema() -> None:
     assert gen_id(FlowSchema({}, ClassE(1).method).flow, [], {}) != gen_id(
         FlowSchema({}, ClassE(2).method).flow, [], {}
     )
+
+
+def _gen_id_calls(_) -> dict[str, str]:
+    gen_ids = {
+        'gen_id(ClassA().method, [], {"a": 1})': gen_id(ClassA().method, [], {"a": 1}),
+        "gen_id(ClassD1().method, [], {})": gen_id(ClassD1().method, [], {}),
+        "gen_id(ClassE(1).method, [], {})": gen_id(ClassE(1).method, [], {}),
+        "gen_id(partial(ClassA().method, 42), [], {})": gen_id(
+            partial(ClassA().method, 42), [], {}
+        ),
+        "gen_id(partial(ClassD1().method, 42), [], {})": gen_id(
+            partial(ClassD1().method, 42), [], {}
+        ),
+        "gen_id(partial(ClassE(1).method, 42), [], {})": gen_id(
+            partial(ClassE(1).method, 42), [], {}
+        ),
+        "gen_id(TaskableA(), [], {})": gen_id(TaskableA(), [], {}),
+        "gen_id(TaskableD(1), [], {})": gen_id(TaskableD(1), [], {}),
+        "gen_id(FlowSchema({}, ClassA().method).flow, [], {})": gen_id(
+            FlowSchema({}, ClassA().method).flow, [], {}
+        ),
+        "gen_id(FlowSchema({}, ClassD1().method).flow, [], {})": gen_id(
+            FlowSchema({}, ClassD1().method).flow, [], {}
+        ),
+        "gen_id(FlowSchema({}, ClassE(1).method).flow, [], {})": gen_id(
+            FlowSchema({}, ClassE(1).method).flow, [], {}
+        ),
+    }
+    return gen_ids
+
+
+def test_persistence_across_sessions() -> None:
+    # Create two separate processes - spawn ensures a new PYTHONHASHSEED is used
+    ctx = multiprocessing.get_context("spawn")
+    with ctx.Pool(processes=2) as pool:
+        result = pool.map(_gen_id_calls, range(2))
+
+    assert result[0] == result[1]