feat(conductor): Implement abort checks in worker lifecycle and kill_worker method
This commit is contained in:
40
tests/test_run_worker_lifecycle_abort.py
Normal file
40
tests/test_run_worker_lifecycle_abort.py
Normal file
@@ -0,0 +1,40 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
import threading
|
||||
import json
|
||||
import time
|
||||
from src.multi_agent_conductor import run_worker_lifecycle
|
||||
from src.models import Ticket, WorkerContext
|
||||
|
||||
class TestRunWorkerLifecycleAbort(unittest.TestCase):
|
||||
def test_run_worker_lifecycle_returns_early_on_abort(self):
|
||||
"""
|
||||
Test that run_worker_lifecycle returns early and marks ticket as 'killed'
|
||||
if the abort event is set for the ticket.
|
||||
"""
|
||||
# Mock ai_client.send
|
||||
with patch('src.ai_client.send') as mock_send:
|
||||
# Mock ticket and context
|
||||
ticket = Ticket(id="T-001", description="Test task")
|
||||
ticket = Ticket(id="T-001", description="Test task")
|
||||
context = WorkerContext(ticket_id="T-001", model_name="test-model")
|
||||
# Mock engine with _abort_events dict
|
||||
mock_engine = MagicMock()
|
||||
abort_event = threading.Event()
|
||||
mock_engine._abort_events = {"T-001": abort_event}
|
||||
|
||||
# Set abort event
|
||||
abort_event.set()
|
||||
|
||||
# Call run_worker_lifecycle
|
||||
# md_content is expected to be passed if called like in ConductorEngine
|
||||
run_worker_lifecycle(ticket, context, engine=mock_engine, md_content="test context")
|
||||
|
||||
# Assert ticket status is 'killed'
|
||||
self.assertEqual(ticket.status, "killed")
|
||||
|
||||
# Also assert ai_client.send was NOT called
|
||||
mock_send.assert_not_called()
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user