fixes

TransformerOptimus · Oct 10, 2023 · 35e8bac · 35e8bac
1 parent cb19f16
commit 35e8bac
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tests/unit_tests/apm/test_call_log_helper.py b/tests/unit_tests/apm/test_call_log_helper.py
@@ -55,13 +55,12 @@ def test_fetch_data_success(call_log_helper, mock_session):
         model='test_model',
         org_id=1
     )]
-    agent = Agent(name='test_agent')
-    tool = Tool(name='test_tool', toolkit_id=1)
-    toolkit = Toolkit(name='test_toolkit')
+    agents = [Agent(name='test_agent')]
+    tools = [Tool(name='test_tool', toolkit_id=1)]
+    toolkits = [Toolkit(name='test_toolkit')]
 
     # setup return values for the mock methods
-    mock_session.query().filter().first.side_effect = [summary_result, agent, tool, toolkit]
-    mock_session.query().filter().all.return_value = runs
+    mock_session.query().filter().first.side_effect = [summary_result, runs, agents, toolkits, tools]
 
     result = call_log_helper.fetch_data('test_model')