Add non-dts tests, update pagination test

Author: andystaples

tests/durabletask-azuremanaged/test_dts_batch_actions.py

@@ -1,4 +1,5 @@
 
+import logging
 import os
 import time
 from datetime import datetime, timedelta, timezone
@@ -140,14 +141,24 @@ def test_get_orchestration_state_by_time_range():
     assert len([o for o in orchestrations_outside_range if o.instance_id == id]) == 0
 
 
-def test_get_orchestration_state_by_max_instance_count():
+def test_get_orchestration_state_pagination_succeeds():
+    # Create a custom handler to capture log messages
+    log_records = []
+
+    class ListHandler(logging.Handler):
+        def emit(self, record):
+            log_records.append(record)
+
+    handler = ListHandler()
+
     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
                                     taskhub=taskhub_name, token_credential=None) as w:
         w.add_orchestrator(empty_orchestrator)
         w.start()
 
         c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
-                                       taskhub=taskhub_name, token_credential=None)
+                                       taskhub=taskhub_name, token_credential=None,
+                                       log_handler=handler)
 
         # Create at least 3 orchestrations to test the limit
         ids = []
@@ -162,8 +173,11 @@ def test_get_orchestration_state_by_max_instance_count():
         # Query with max_instance_count=2
         orchestrations = c.get_orchestration_state_by(max_instance_count=2)
 
-    # Should return exactly 2 instances since we created at least 3
-    assert len(orchestrations) == 2
+    # Should return more than 2 instances since we created at least 3
+    assert len(orchestrations) > 2
+    # Verify the pagination loop ran by checking for the continuation token log message
+    assert any("Received continuation token" in record.getMessage() for record in log_records), \
+        "Expected pagination loop to execute with continuation token"
 
 
 def test_purge_orchestration():

tests/durabletask/test_batch_actions.py

@@ -1,4 +1,5 @@
 
+from datetime import datetime, timedelta, timezone
 from grpc._channel import _InactiveRpcError
 import pytest
 from durabletask import client, task, worker
@@ -13,6 +14,10 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
     return "Complete"
 
 
+def failing_orchestrator(ctx: task.OrchestrationContext, _):
+    raise Exception("Orchestration failed")
+
+
 def test_get_all_orchestration_states():
     # Start a worker, which will connect to the sidecar in a background thread
     with worker.TaskHubGrpcWorker() as w:
@@ -26,3 +31,96 @@ def test_get_all_orchestration_states():
         with pytest.raises(_InactiveRpcError) as exec_info:
             c.get_all_orchestration_states()
         assert "unimplemented" in str(exec_info.value)
+
+
+def test_get_all_entities():
+    # Start a worker, which will connect to the sidecar in a background thread
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.start()
+
+        c = client.TaskHubGrpcClient()
+        with pytest.raises(_InactiveRpcError) as exec_info:
+            c.get_all_entities()
+        assert "unimplemented" in str(exec_info.value)
+
+
+def test_clean_entity_storage():
+    # Start a worker, which will connect to the sidecar in a background thread
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.start()
+
+        c = client.TaskHubGrpcClient()
+        with pytest.raises(_InactiveRpcError) as exec_info:
+            c.clean_entity_storage()
+        assert "unimplemented" in str(exec_info.value)
+
+
+def test_purge_orchestrations_by_status():
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(failing_orchestrator)
+        w.start()
+
+        c = client.TaskHubGrpcClient()
+
+        # Schedule and let it fail
+        failed_id = c.schedule_new_orchestration(failing_orchestrator)
+        try:
+            c.wait_for_orchestration_completion(failed_id, timeout=30)
+        except client.OrchestrationFailedError:
+            pass  # Expected failure
+
+        # Verify it exists and is failed
+        state_before = c.get_orchestration_state(failed_id)
+        assert state_before is not None
+        assert state_before.runtime_status == client.OrchestrationStatus.FAILED
+
+        # Purge failed orchestrations
+        result = c.purge_orchestrations_by(
+            runtime_status=[client.OrchestrationStatus.FAILED],
+            recursive=True
+        )
+
+        # Verify purge result
+        assert result.deleted_instance_count >= 1
+
+        # Verify the failed orchestration no longer exists
+        state_after = c.get_orchestration_state(failed_id)
+        assert state_after is None
+
+
+def test_purge_orchestrations_by_time_range():
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.start()
+
+        c = client.TaskHubGrpcClient()
+
+        # Get current time
+        before_creation = datetime.now(timezone.utc) - timedelta(seconds=5)
+
+        # Schedule orchestration
+        id = c.schedule_new_orchestration(empty_orchestrator, input="ToPurgeByTime")
+        c.wait_for_orchestration_completion(id, timeout=30)
+
+        after_creation = datetime.now(timezone.utc) + timedelta(seconds=5)
+
+        # Verify it exists
+        state_before = c.get_orchestration_state(id)
+        assert state_before is not None
+
+        # Purge by time range
+        result = c.purge_orchestrations_by(
+            created_time_from=before_creation,
+            created_time_to=after_creation,
+            runtime_status=[client.OrchestrationStatus.COMPLETED],
+            recursive=True
+        )
+
+        # Verify purge result
+        assert result.deleted_instance_count >= 1
+
+        # Verify it no longer exists
+        state_after = c.get_orchestration_state(id)
+        assert state_after is None