more test imports

Author: berndverst

tests/durabletask/test_worker_concurrency_loop.py

@@ -2,7 +2,7 @@
 import threading
 import time
 
-from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
+from durabletask import worker
 
 
 class DummyStub:
@@ -37,12 +37,12 @@ class DummyCompletionToken:
     pass
 
 def test_worker_concurrency_loop_sync():
-    options = ConcurrencyOptions(
+    options = worker.ConcurrencyOptions(
         maximum_concurrent_activity_work_items=2,
         maximum_concurrent_orchestration_work_items=1,
         maximum_thread_pool_workers=2,
     )
-    worker = TaskHubGrpcWorker(concurrency_options=options)
+    grpc_worker = worker.TaskHubGrpcWorker(concurrency_options=options)
     stub = DummyStub()
 
     def dummy_orchestrator(req, stub, completionToken):
@@ -54,25 +54,25 @@ def dummy_activity(req, stub, completionToken):
         stub.CompleteActivityTask('ok')
 
     # Patch the worker's _execute_orchestrator and _execute_activity
-    worker._execute_orchestrator = dummy_orchestrator
-    worker._execute_activity = dummy_activity
+    grpc_worker._execute_orchestrator = dummy_orchestrator
+    grpc_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(worker._async_worker_manager.run())
+        worker_task = asyncio.create_task(grpc_worker._async_worker_manager.run())
         for req in orchestrator_requests:
-            worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
+            grpc_worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
         for req in activity_requests:
-            worker._async_worker_manager.submit_activity(dummy_activity, req, stub, DummyCompletionToken())
+            grpc_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}"
-        worker._async_worker_manager._shutdown = True
+        grpc_worker._async_worker_manager._shutdown = True
         await worker_task
     asyncio.run(run_test())
 

tests/durabletask/test_worker_concurrency_loop_async.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
+from durabletask import worker
 
 
 class DummyStub:
@@ -38,12 +38,12 @@ class DummyCompletionToken:
 
 
 def test_worker_concurrency_loop_async():
-    options = ConcurrencyOptions(
+    options = worker.ConcurrencyOptions(
         maximum_concurrent_activity_work_items=2,
         maximum_concurrent_orchestration_work_items=1,
         maximum_thread_pool_workers=2,
     )
-    worker = TaskHubGrpcWorker(concurrency_options=options)
+    grpc_worker = worker.TaskHubGrpcWorker(concurrency_options=options)
     stub = DummyStub()
 
     async def dummy_orchestrator(req, stub, completionToken):
@@ -55,26 +55,26 @@ async def dummy_activity(req, stub, completionToken):
         stub.CompleteActivityTask('ok')
 
     # Patch the worker's _execute_orchestrator and _execute_activity
-    worker._execute_orchestrator = dummy_orchestrator
-    worker._execute_activity = dummy_activity
+    grpc_worker._execute_orchestrator = dummy_orchestrator
+    grpc_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():
         # Clear stub state before each run
         stub.completed.clear()
-        worker_task = asyncio.create_task(worker._async_worker_manager.run())
+        worker_task = asyncio.create_task(grpc_worker._async_worker_manager.run())
         for req in orchestrator_requests:
-            worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
+            grpc_worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
         for req in activity_requests:
-            worker._async_worker_manager.submit_activity(dummy_activity, req, stub, DummyCompletionToken())
+            grpc_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}"
-        worker._async_worker_manager._shutdown = True
+        grpc_worker._async_worker_manager._shutdown = True
         await worker_task
     asyncio.run(run_test())
     asyncio.run(run_test())