fix: two more tests

Author: ringoldsdev

laygo/context/parallel.py

@@ -7,6 +7,7 @@
 from collections.abc import Iterator
 import multiprocessing as mp
 from multiprocessing.managers import DictProxy
+import threading
 from threading import Lock
 from typing import Any
 from typing import TypeVar
@@ -64,23 +65,24 @@ def __init__(self, initial_context: dict[str, Any] | None = None, handle: Parall
       self._shared_dict = self._manager.dict(initial_context or {})
       self._lock = self._manager.Lock()
 
-    self._is_locked = False
+    # Thread-local storage for lock state to handle concurrent access
+    self._local = threading.local()
 
   def _lock_context(self) -> None:
     """Acquire the lock for this context manager."""
-    if not self._is_locked:
+    if not getattr(self._local, "is_locked", False):
       self._lock.acquire()
-      self._is_locked = True
+      self._local.is_locked = True
 
   def _unlock_context(self) -> None:
     """Release the lock for this context manager."""
-    if self._is_locked:
+    if getattr(self._local, "is_locked", False):
       self._lock.release()
-      self._is_locked = False
+      self._local.is_locked = False
 
   def _execute_locked(self, operation: Callable[[], R]) -> R:
     """A private helper to execute an operation within a lock."""
-    if not self._is_locked:
+    if not getattr(self._local, "is_locked", False):
       self._lock_context()
       try:
         return operation()

tests/test_integration.py

@@ -134,13 +134,13 @@ def test_parallel_transformer_with_context_modification(self):
     parallel_transformer = parallel_transformer.map(safe_increment_and_transform)
 
     data = [1, 2, 3, 4, 5]
-    result, _ = Pipeline(data).context(context).apply(parallel_transformer).to_list()
+    result, processed_context = Pipeline(data).context(context).apply(parallel_transformer).to_list()
 
     # Verify transformation results
     assert sorted(result) == [2, 4, 6, 8, 10]
     # Verify context was safely modified
-    assert context["processed_count"] == len(data)
-    assert context["sum_total"] == sum(data)
+    assert processed_context["processed_count"] == len(data)
+    assert processed_context["sum_total"] == sum(data)
 
   def test_pipeline_accesses_modified_context(self):
     """Test that pipeline can access context data modified by parallel transformer."""

tests/test_threaded_transformer.py

@@ -100,9 +100,10 @@ def test_context_modification_with_locking(self):
     context = ParallelContextManager({"items": 0})
 
     def safe_increment(x: int, ctx: IContextManager) -> int:
-      current_items = ctx["items"]
-      time.sleep(0.001)  # Increase chance of race condition
-      ctx["items"] = current_items + 1
+      with ctx:
+        # Simulate a race condition
+        time.sleep(0.001)  # Increase chance of race condition
+        ctx["items"] = ctx["items"] + 1
       return x * 2
 
     transformer = ThreadedTransformer[int, int](max_workers=4, chunk_size=1)