PR Feedback

Author: andystaples

durabletask/testing/in_memory_backend.py

@@ -160,7 +160,9 @@ def stop(self, grace: Optional[float] = None):
         self._shutdown_event.set()
         self._work_available.set()  # Unblock GetWorkItems loops
         if self._server:
-            self._server.stop(grace)
+            stop_future = self._server.stop(grace)
+            stop_future.wait()
+            self._server = None
             self._logger.info("In-memory backend stopped")
 
     def reset(self):
@@ -621,6 +623,8 @@ def CompleteEntityTask(self, request: pb.EntityBatchResult, context):
             # Update entity state
             if request.entityState and request.entityState.value:
                 entity.serialized_state = request.entityState.value
+            else:
+                entity.serialized_state = None
             entity.last_modified_at = datetime.now(timezone.utc)
 
             # Update completion token for next batch
@@ -1472,7 +1476,8 @@ def create_test_backend(port: int = 50051, max_history_size: int = 10000) -> InM
         ```python
         import pytest
         from durabletask.testing.in_memory_backend import create_test_backend
-        from durabletask import TaskHubGrpcClient, TaskHubGrpcWorker
+        from durabletask.client import TaskHubGrpcClient
+        from durabletask.worker import TaskHubGrpcWorker
 
         @pytest.fixture
         def backend():

tests/durabletask/entities/test_class_based_entities_e2e.py

@@ -12,13 +12,13 @@
 from durabletask import client, entities, task, worker
 from durabletask.testing import create_test_backend
 
-HOST = "localhost:50055"
+HOST = "localhost:50059"
 
 
 @pytest.fixture(autouse=True)
 def backend():
     """Create an in-memory backend for entity testing."""
-    b = create_test_backend(port=50055)
+    b = create_test_backend(port=50059)
     yield b
     b.stop()
     b.reset()

tests/durabletask/test_orchestration_versioning_e2e.py

@@ -2,7 +2,6 @@
 # Licensed under the MIT License.
 
 import json
-import warnings
 
 import pytest
 
@@ -22,9 +21,6 @@ def backend():
 
 
 def test_versioned_orchestration_succeeds():
-    warnings.warn("Skipping test_versioned_orchestration_succeeds. "
-                  "Currently not passing as the in-memory backend does not support versioning yet")
-
     def plus_one(_: task.ActivityContext, input: int) -> int:
         return input + 1