Add support to RestartInstance to backend, unskip tests

Author: andystaples

durabletask/testing/in_memory_backend.py

@@ -395,6 +395,40 @@ def PurgeInstances(self, request: pb.PurgeInstancesRequest, context):
             isComplete=wrappers_pb2.BoolValue(value=True),
         )
 
+    def RestartInstance(self, request: pb.RestartInstanceRequest, context):
+        """Restarts a completed orchestration instance."""
+        with self._lock:
+            instance = self._instances.get(request.instanceId)
+            if not instance:
+                context.abort(
+                    grpc.StatusCode.NOT_FOUND,
+                    f"Orchestration instance '{request.instanceId}' not found")
+                return pb.RestartInstanceResponse()
+
+            if not self._is_terminal_status(instance.status):
+                context.abort(
+                    grpc.StatusCode.FAILED_PRECONDITION,
+                    f"Orchestration instance '{request.instanceId}' is not in a terminal state")
+                return pb.RestartInstanceResponse()
+
+            name = instance.name
+            original_input = instance.input
+
+            if request.restartWithNewInstanceId:
+                new_instance_id = uuid.uuid4().hex
+            else:
+                new_instance_id = request.instanceId
+                # Remove the old instance so we can recreate it
+                del self._instances[request.instanceId]
+                self._orchestration_queue_set.discard(request.instanceId)
+                self._state_waiters.pop(request.instanceId, None)
+
+            self._create_instance_internal(new_instance_id, name, original_input)
+
+        self._logger.info(
+            f"Restarted instance '{request.instanceId}' as '{new_instance_id}'")
+        return pb.RestartInstanceResponse(instanceId=new_instance_id)
+
     def GetWorkItems(self, request: pb.GetWorkItemsRequest, context):
         """Streams work items to the worker (orchestration and activity work items)."""
         self._logger.info("Worker connected and requesting work items")

tests/durabletask/test_orchestration_e2e.py

@@ -343,7 +343,6 @@ def child(ctx: task.OrchestrationContext, _):
         assert state is None
 
 
-@pytest.mark.skip(reason="durabletask-go does not yet support RestartInstance")
 def test_restart_with_same_instance_id():
     def orchestrator(ctx: task.OrchestrationContext, _):
         result = yield ctx.call_activity(say_hello, input="World")
@@ -353,12 +352,12 @@ def say_hello(ctx: task.ActivityContext, input: str):
         return f"Hello, {input}!"
 
     # Start a worker, which will connect to the sidecar in a background thread
-    with worker.TaskHubGrpcWorker() as w:
+    with worker.TaskHubGrpcWorker(host_address=HOST) as w:
         w.add_orchestrator(orchestrator)
         w.add_activity(say_hello)
         w.start()
 
-        task_hub_client = client.TaskHubGrpcClient()
+        task_hub_client = client.TaskHubGrpcClient(host_address=HOST)
         id = task_hub_client.schedule_new_orchestration(orchestrator)
         state = task_hub_client.wait_for_orchestration_completion(id, timeout=30)
         assert state is not None
@@ -375,7 +374,6 @@ def say_hello(ctx: task.ActivityContext, input: str):
         assert state.serialized_output == json.dumps("Hello, World!")
 
 
-@pytest.mark.skip(reason="durabletask-go does not yet support RestartInstance")
 def test_restart_with_new_instance_id():
     def orchestrator(ctx: task.OrchestrationContext, _):
         result = yield ctx.call_activity(say_hello, input="World")
@@ -385,12 +383,12 @@ def say_hello(ctx: task.ActivityContext, input: str):
         return f"Hello, {input}!"
 
     # Start a worker, which will connect to the sidecar in a background thread
-    with worker.TaskHubGrpcWorker() as w:
+    with worker.TaskHubGrpcWorker(host_address=HOST) as w:
         w.add_orchestrator(orchestrator)
         w.add_activity(say_hello)
         w.start()
 
-        task_hub_client = client.TaskHubGrpcClient()
+        task_hub_client = client.TaskHubGrpcClient(host_address=HOST)
         id = task_hub_client.schedule_new_orchestration(orchestrator)
         state = task_hub_client.wait_for_orchestration_completion(id, timeout=30)
         assert state is not None