increase test coverage on new code

Signed-off-by: Filinto Duran <[email protected]>

Author: filintod

tests/durabletask/test_client.py

@@ -1,4 +1,4 @@
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
 from durabletask.internal.shared import get_default_host_address, get_grpc_channel
@@ -140,3 +140,17 @@ def test_sync_channel_passes_base_options_and_max_lengths():
         assert ("grpc.max_send_message_length", 1234) in opts
         assert ("grpc.max_receive_message_length", 5678) in opts
         assert ("grpc.primary_user_agent", "durabletask-tests") in opts
+
+
+def test_taskhub_client_close_handles_exceptions():
+    """Test that close() handles exceptions gracefully (edge case not easily testable in E2E)."""
+    with patch("durabletask.internal.shared.get_grpc_channel") as mock_get_channel:
+        mock_channel = MagicMock()
+        mock_channel.close.side_effect = Exception("close failed")
+        mock_get_channel.return_value = mock_channel
+
+        from durabletask import client
+
+        task_hub_client = client.TaskHubGrpcClient()
+        # Should not raise exception
+        task_hub_client.close()

tests/durabletask/test_orchestration_e2e.py

@@ -65,6 +65,11 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
         id = c.schedule_new_orchestration(empty_orchestrator)
         state = c.wait_for_orchestration_completion(id, timeout=30)
 
+        # Test calling wait again on already-completed orchestration (should return immediately)
+        state2 = c.wait_for_orchestration_completion(id, timeout=30)
+        assert state2 is not None
+        assert state2.runtime_status == client.OrchestrationStatus.COMPLETED
+
     assert invoked
     assert state is not None
     assert state.name == task.get_name(empty_orchestrator)
@@ -76,6 +81,41 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
     assert state.serialized_custom_status is None
 
 
+def test_wait_for_idle():
+    """Test that wait_for_idle properly waits for in-flight activities to complete."""
+    import time
+
+    def slow_activity(ctx: task.ActivityContext, input: int):
+        # Simulate slow activity
+        time.sleep(0.2)
+        return input + 1
+
+    def orchestrator(ctx: task.OrchestrationContext, input: int):
+        # Schedule multiple activities without waiting
+        tasks = [ctx.call_activity(slow_activity, input=i) for i in range(3)]
+        # Wait for all to complete
+        results = yield task.when_all(tasks)
+        return sum(results)
+
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(orchestrator)
+        w.add_activity(slow_activity)
+        w.start()
+        w.wait_for_ready(timeout=10)
+
+        with client.TaskHubGrpcClient() as c:
+            id = c.schedule_new_orchestration(orchestrator, input=1)
+
+            # Wait for orchestration to complete
+            state = c.wait_for_orchestration_completion(id, timeout=30)
+            assert state is not None
+            assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+
+        # Wait for any lingering activities to finish
+        idle = w.wait_for_idle(timeout=5.0)
+        assert idle is True
+
+
 def test_activity_sequence():
     def plus_one(_: task.ActivityContext, input: int) -> int:
         return input + 1

tests/durabletask/test_worker_grpc_errors.py

@@ -0,0 +1,114 @@
+"""
+Copyright 2025 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from unittest.mock import MagicMock, Mock, patch
+
+import grpc
+
+from durabletask import worker
+
+
+def test_execute_orchestrator_grpc_error_benign_cancelled():
+    """Test that benign gRPC errors in orchestrator execution are handled gracefully."""
+    w = worker.TaskHubGrpcWorker()
+
+    # Add a dummy orchestrator
+    def test_orchestrator(ctx, input):
+        return "result"
+
+    w.add_orchestrator(test_orchestrator)
+
+    # Mock the stub to raise a benign error
+    mock_stub = MagicMock()
+    mock_error = grpc.RpcError()
+    mock_error.code = Mock(return_value=grpc.StatusCode.CANCELLED)
+    mock_stub.CompleteOrchestratorTask.side_effect = mock_error
+
+    # Create a mock request with proper structure
+    mock_req = MagicMock()
+    mock_req.instanceId = "test-id"
+    mock_req.pastEvents = []
+    mock_req.newEvents = [MagicMock()]
+    mock_req.newEvents[0].HasField = lambda x: x == "executionStarted"
+    mock_req.newEvents[0].executionStarted.name = "test_orchestrator"
+    mock_req.newEvents[0].executionStarted.input = None
+    mock_req.newEvents[0].router.targetAppID = None
+    mock_req.newEvents[0].router.sourceAppID = None
+    mock_req.newEvents[0].timestamp.ToDatetime = Mock(return_value=None)
+
+    # Should not raise exception (benign error)
+    w._execute_orchestrator(mock_req, mock_stub, "token")
+
+
+def test_execute_orchestrator_grpc_error_non_benign():
+    """Test that non-benign gRPC errors in orchestrator execution are logged."""
+    w = worker.TaskHubGrpcWorker()
+
+    # Add a dummy orchestrator
+    def test_orchestrator(ctx, input):
+        return "result"
+
+    w.add_orchestrator(test_orchestrator)
+
+    # Mock the stub to raise a non-benign error
+    mock_stub = MagicMock()
+    mock_error = grpc.RpcError()
+    mock_error.code = Mock(return_value=grpc.StatusCode.INTERNAL)
+    mock_stub.CompleteOrchestratorTask.side_effect = mock_error
+
+    # Create a mock request with proper structure
+    mock_req = MagicMock()
+    mock_req.instanceId = "test-id"
+    mock_req.pastEvents = []
+    mock_req.newEvents = [MagicMock()]
+    mock_req.newEvents[0].HasField = lambda x: x == "executionStarted"
+    mock_req.newEvents[0].executionStarted.name = "test_orchestrator"
+    mock_req.newEvents[0].executionStarted.input = None
+    mock_req.newEvents[0].router.targetAppID = None
+    mock_req.newEvents[0].router.sourceAppID = None
+    mock_req.newEvents[0].timestamp.ToDatetime = Mock(return_value=None)
+
+    # Should not raise exception (error is logged but handled)
+    with patch.object(w._logger, "exception") as mock_log:
+        w._execute_orchestrator(mock_req, mock_stub, "token")
+        # Verify error was logged
+        assert mock_log.called
+
+
+def test_execute_activity_grpc_error_benign():
+    """Test that benign gRPC errors in activity execution are handled gracefully."""
+    w = worker.TaskHubGrpcWorker()
+
+    # Add a dummy activity
+    def test_activity(ctx, input):
+        return "result"
+
+    w.add_activity(test_activity)
+
+    # Mock the stub to raise a benign error
+    mock_stub = MagicMock()
+    mock_error = grpc.RpcError()
+    mock_error.code = Mock(return_value=grpc.StatusCode.CANCELLED)
+    str_return = "unknown instance ID/task ID combo"
+    mock_error.__str__ = Mock(return_value=str_return)
+    mock_stub.CompleteActivityTask.side_effect = mock_error
+
+    # Create a mock request
+    mock_req = MagicMock()
+    mock_req.orchestrationInstance.instanceId = "test-id"
+    mock_req.name = "test_activity"
+    mock_req.taskId = 1
+    mock_req.input.value = '""'
+
+    # Should not raise exception (benign error)
+    w._execute_activity(mock_req, mock_stub, "token")