From 7a2c7668245ba353f772fbd94d190b85ceb68553 Mon Sep 17 00:00:00 2001 From: Ariel Frischer Date: Thu, 13 Feb 2025 20:18:02 -0800 Subject: [PATCH] test(agent_utils): add config parameter to mock_ciayn assertions for better clarity and flexibility in tests --- tests/ra_aid/test_agent_utils.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/ra_aid/test_agent_utils.py b/tests/ra_aid/test_agent_utils.py index 726c7d7..082c140 100644 --- a/tests/ra_aid/test_agent_utils.py +++ b/tests/ra_aid/test_agent_utils.py @@ -128,7 +128,9 @@ def test_create_agent_openai(mock_model, mock_memory): assert agent == "ciayn_agent" mock_ciayn.assert_called_once_with( - mock_model, [], max_tokens=models_params["openai"]["gpt-4"]["token_limit"] + mock_model, [], + max_tokens=models_params["openai"]["gpt-4"]["token_limit"], + config={'provider': 'openai', 'model': 'gpt-4'} ) @@ -142,7 +144,9 @@ def test_create_agent_no_token_limit(mock_model, mock_memory): assert agent == "ciayn_agent" mock_ciayn.assert_called_once_with( - mock_model, [], max_tokens=DEFAULT_TOKEN_LIMIT + mock_model, [], + max_tokens=DEFAULT_TOKEN_LIMIT, + config={'provider': 'unknown', 'model': 'unknown-model'} ) @@ -159,6 +163,7 @@ def test_create_agent_missing_config(mock_model, mock_memory): mock_model, [], max_tokens=DEFAULT_TOKEN_LIMIT, + config={'provider': 'openai'} ) @@ -202,7 +207,9 @@ def test_create_agent_with_checkpointer(mock_model, mock_memory): assert agent == "ciayn_agent" mock_ciayn.assert_called_once_with( - mock_model, [], max_tokens=models_params["openai"]["gpt-4"]["token_limit"] + mock_model, [], + max_tokens=models_params["openai"]["gpt-4"]["token_limit"], + config={'provider': 'openai', 'model': 'gpt-4'} )