cleanup tests

Author: berndverst

tests/durabletask/test_concurrency_options.py

@@ -3,12 +3,12 @@
 
 import os
 
-from durabletask import worker
+from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
 
 def test_default_concurrency_options():
     """Test that default concurrency options work correctly."""
-    options = worker.ConcurrencyOptions()
+    options = ConcurrencyOptions()
     processor_count = os.cpu_count() or 1
     expected_default = 100 * processor_count
     expected_workers = processor_count + 4
@@ -20,7 +20,7 @@ def test_default_concurrency_options():
 
 def test_custom_concurrency_options():
     """Test that custom concurrency options work correctly."""
-    options = worker.ConcurrencyOptions(
+    options = ConcurrencyOptions(
         maximum_concurrent_activity_work_items=50,
         maximum_concurrent_orchestration_work_items=25,
         maximum_thread_pool_workers=30,
@@ -37,7 +37,7 @@ def test_partial_custom_options():
     expected_default = 100 * processor_count
     expected_workers = processor_count + 4
 
-    options = worker.ConcurrencyOptions(
+    options = ConcurrencyOptions(
         maximum_concurrent_activity_work_items=30
     )
 
@@ -48,44 +48,44 @@ def test_partial_custom_options():
 
 def test_worker_with_concurrency_options():
     """Test that TaskHubGrpcWorker accepts concurrency options."""
-    options = worker.ConcurrencyOptions(
+    options = ConcurrencyOptions(
         maximum_concurrent_activity_work_items=10,
         maximum_concurrent_orchestration_work_items=20,
         maximum_thread_pool_workers=15,
     )
 
-    grpc_worker = worker.TaskHubGrpcWorker(concurrency_options=options)
+    worker = TaskHubGrpcWorker(concurrency_options=options)
 
-    assert grpc_worker.concurrency_options == options
+    assert worker.concurrency_options == options
 
 
 def test_worker_default_options():
     """Test that TaskHubGrpcWorker uses default options when no parameters are provided."""
-    grpc_worker = worker.TaskHubGrpcWorker()
+    worker = TaskHubGrpcWorker()
 
     processor_count = os.cpu_count() or 1
     expected_default = 100 * processor_count
     expected_workers = processor_count + 4
 
     assert (
-        grpc_worker.concurrency_options.maximum_concurrent_activity_work_items == expected_default
+        worker.concurrency_options.maximum_concurrent_activity_work_items == expected_default
     )
     assert (
-        grpc_worker.concurrency_options.maximum_concurrent_orchestration_work_items == expected_default
+        worker.concurrency_options.maximum_concurrent_orchestration_work_items == expected_default
     )
-    assert grpc_worker.concurrency_options.maximum_thread_pool_workers == expected_workers
+    assert worker.concurrency_options.maximum_thread_pool_workers == expected_workers
 
 
 def test_concurrency_options_property_access():
     """Test that the concurrency_options property works correctly."""
-    options = worker.ConcurrencyOptions(
+    options = ConcurrencyOptions(
         maximum_concurrent_activity_work_items=15,
         maximum_concurrent_orchestration_work_items=25,
         maximum_thread_pool_workers=30,
     )
 
-    grpc_worker = worker.TaskHubGrpcWorker(concurrency_options=options)
-    retrieved_options = grpc_worker.concurrency_options
+    worker = TaskHubGrpcWorker(concurrency_options=options)
+    retrieved_options = worker.concurrency_options
 
     # Should be the same object
     assert retrieved_options is options

tests/durabletask/test_worker_concurrency_loop.py

@@ -2,7 +2,7 @@
 import threading
 import time
 
-from durabletask import worker
+from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
 
 class DummyStub:
@@ -15,6 +15,7 @@ def CompleteOrchestratorTask(self, res):
     def CompleteActivityTask(self, res):
         self.completed.append(('activity', res))
 
+
 class DummyRequest:
     def __init__(self, kind, instance_id):
         self.kind = kind
@@ -33,16 +34,18 @@ def HasField(self, field):
     def WhichOneof(self, _):
         return f'{self.kind}Request'
 
+
 class DummyCompletionToken:
     pass
 
+
 def test_worker_concurrency_loop_sync():
-    options = worker.ConcurrencyOptions(
+    options = ConcurrencyOptions(
         maximum_concurrent_activity_work_items=2,
         maximum_concurrent_orchestration_work_items=1,
         maximum_thread_pool_workers=2,
     )
-    grpc_worker = worker.TaskHubGrpcWorker(concurrency_options=options)
+    worker = TaskHubGrpcWorker(concurrency_options=options)
     stub = DummyStub()
 
     def dummy_orchestrator(req, stub, completionToken):
@@ -54,39 +57,42 @@ def dummy_activity(req, stub, completionToken):
         stub.CompleteActivityTask('ok')
 
     # Patch the worker's _execute_orchestrator and _execute_activity
-    grpc_worker._execute_orchestrator = dummy_orchestrator
-    grpc_worker._execute_activity = dummy_activity
+    worker._execute_orchestrator = dummy_orchestrator
+    worker._execute_activity = dummy_activity
 
     orchestrator_requests = [DummyRequest('orchestrator', f'orch{i}') for i in range(3)]
     activity_requests = [DummyRequest('activity', f'act{i}') for i in range(4)]
 
     async def run_test():
         # Start the worker manager's run loop in the background
-        worker_task = asyncio.create_task(grpc_worker._async_worker_manager.run())
+        worker_task = asyncio.create_task(worker._async_worker_manager.run())
         for req in orchestrator_requests:
-            grpc_worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
+            worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
         for req in activity_requests:
-            grpc_worker._async_worker_manager.submit_activity(dummy_activity, req, stub, DummyCompletionToken())
+            worker._async_worker_manager.submit_activity(dummy_activity, req, stub, DummyCompletionToken())
         await asyncio.sleep(1.0)
         orchestrator_count = sum(1 for t, _ in stub.completed if t == 'orchestrator')
         activity_count = sum(1 for t, _ in stub.completed if t == 'activity')
         assert orchestrator_count == 3, f"Expected 3 orchestrator completions, got {orchestrator_count}"
         assert activity_count == 4, f"Expected 4 activity completions, got {activity_count}"
-        grpc_worker._async_worker_manager._shutdown = True
+        worker._async_worker_manager._shutdown = True
         await worker_task
     asyncio.run(run_test())
 
+
 # Dummy orchestrator and activity for sync context
 def dummy_orchestrator(ctx, input):
     # Simulate some work
     time.sleep(0.1)
     return "orchestrator-done"
 
+
 def dummy_activity(ctx, input):
     # Simulate some work
     time.sleep(0.1)
     return "activity-done"
 
+
 def test_worker_concurrency_sync():
     # Use small concurrency to make test observable
     options = ConcurrencyOptions(
@@ -124,16 +130,11 @@ def run_manager():
     t = threading.Thread(target=run_manager)
     t.start()
     time.sleep(1.5)  # Let work process
-
-    # Signal shutdown but don't actually call shutdown() yet
-    manager._shutdown = True
+    manager.shutdown()
     # Unblock the consumers by putting dummy items in the queues
     manager.activity_queue.put_nowait((lambda: None, (), {}))
     manager.orchestration_queue.put_nowait((lambda: None, (), {}))
     t.join(timeout=2)
 
-    # Now shutdown the thread pool
-    manager.thread_pool.shutdown(wait=True)
-
     # Check that all work items completed
-    assert len(results) == 10
+    assert len(results) == 10

tests/durabletask/test_worker_concurrency_loop_async.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from durabletask import worker
+from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
 
 class DummyStub:
@@ -38,12 +38,12 @@ class DummyCompletionToken:
 
 
 def test_worker_concurrency_loop_async():
-    options = worker.ConcurrencyOptions(
+    options = ConcurrencyOptions(
         maximum_concurrent_activity_work_items=2,
         maximum_concurrent_orchestration_work_items=1,
         maximum_thread_pool_workers=2,
     )
-    grpc_worker = worker.TaskHubGrpcWorker(concurrency_options=options)
+    grpc_worker = TaskHubGrpcWorker(concurrency_options=options)
     stub = DummyStub()
 
     async def dummy_orchestrator(req, stub, completionToken):