diff --git a/tests/test_orchestration_logic.py b/tests/test_orchestration_logic.py index 1ff44224..c4948d79 100644 --- a/tests/test_orchestration_logic.py +++ b/tests/test_orchestration_logic.py @@ -4,6 +4,7 @@ from src import orchestrator_pm from src import multi_agent_conductor from src import conductor_tech_lead from src.models import Ticket, Track, WorkerContext +from src.result_types import Result def test_generate_tracks() -> None: mock_response = """ @@ -12,7 +13,7 @@ def test_generate_tracks() -> None: {"id": "track_2", "title": "Refactor", "goal": "decouple modules", "type": "refactor"} ] """ - with patch("src.ai_client.send", return_value=mock_response): + with patch("src.ai_client.send_result", return_value=Result(data=mock_response)): tracks = orchestrator_pm.generate_tracks("Develop feature X", {}, []) assert len(tracks) == 2 assert tracks[0]["id"] == "track_1" @@ -25,7 +26,7 @@ def test_generate_tickets() -> None: {"id": "T2", "description": "task 2", "depends_on": ["T1"]} ] """ - with patch("src.ai_client.send", return_value=mock_response): + with patch("src.ai_client.send_result", return_value=Result(data=mock_response)): tickets = conductor_tech_lead.generate_tickets("Track goal", "code skeletons") assert len(tickets) == 2 assert tickets[0]["id"] == "T1"