From 9feb981ca69b9b0387a4cdfb25e5f868b79e270a Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Mon, 26 Aug 2024 18:48:45 +0300 Subject: [PATCH] fix: Fix execute yaml task test --- agents-api/tests/test_execution_workflow.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/agents-api/tests/test_execution_workflow.py b/agents-api/tests/test_execution_workflow.py index 395f77de9..2dd79e017 100644 --- a/agents-api/tests/test_execution_workflow.py +++ b/agents-api/tests/test_execution_workflow.py @@ -719,7 +719,7 @@ async def _( ): mock_model_response = ModelResponse( id="fake_id", - choices=[Choices(message={"role": "assistant", "content": "Hello, world!"})], + choices=[Choices(message={"role": "assistant", "content": "found: true\nvalue: 'Gaga'"})], created=0, object="text_completion", ) @@ -727,9 +727,14 @@ async def _( with patch("agents_api.clients.litellm.acompletion") as acompletion, open( "./tests/sample_tasks/find_selector.yaml", "r" ) as task_file: + input = dict( + screenshot_base64="iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAACXBIWXMAAAsTAAALEwEAmpwYAAAA", + network_requests=[{"request": {}, "response": {"body": "Lady Gaga"}}], + parameters=["name"], + ) task_definition = yaml.safe_load(task_file) acompletion.return_value = mock_model_response - data = CreateExecutionRequest(input={"parameters": ["param1", "param2"]}) + data = CreateExecutionRequest(input=input) task = create_task( developer_id=developer_id, @@ -752,6 +757,4 @@ async def _( mock_run_task_execution_workflow.assert_called_once() - result = await handle.result() - assert result["content"] == "Hello, world!" - assert result["role"] == "assistant" + await handle.result()