From b6152f805e931f15d562f2fb6f30f0998bf4b1c8 Mon Sep 17 00:00:00 2001 From: Joshua Carroll Date: Thu, 8 Aug 2024 20:23:29 -0700 Subject: [PATCH] Fix broken test from agent import change --- service/test_service.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/service/test_service.py b/service/test_service.py index 974c923..06d9fb5 100644 --- a/service/test_service.py +++ b/service/test_service.py @@ -7,20 +7,19 @@ client = TestClient(app) -@patch("service.service.construct_agent") -def test_invoke(mock_construct_agent): - agent = mock_construct_agent.return_value +@patch("service.service.research_assistant") +def test_invoke(mock_agent): QUESTION = "What is the weather in Tokyo?" ANSWER = "The weather in Tokyo is 70 degrees." agent_response = {"messages": [AIMessage(content=ANSWER)]} - agent.ainvoke = AsyncMock(return_value=agent_response) + mock_agent.ainvoke = AsyncMock(return_value=agent_response) with client as c: response = c.post("/invoke", json={"message": QUESTION}) assert response.status_code == 200 - agent.ainvoke.assert_awaited_once() - input_message = agent.ainvoke.await_args.kwargs["input"]["messages"][0] + mock_agent.ainvoke.assert_awaited_once() + input_message = mock_agent.ainvoke.await_args.kwargs["input"]["messages"][0] assert input_message.content == QUESTION output = ChatMessage.parse_obj(response.json())