hacking

Author: donald-pinckney

tests/openai_agents/basic/test_hello_world_workflow.py

@@ -8,34 +8,43 @@
 from temporalio.worker import Worker
 
 from openai_agents.basic.workflows.hello_world_workflow import HelloWorldAgent
+from tests.openai_agents.conftest import sequential_test_model
 
 
 @pytest.fixture
-def mocked_model(mocker):
-    mock = mocker.AsyncMock()
-    mock.get_response.side_effect = [
-        ModelResponse(
-            output=[
-                ResponseOutputMessage(
-                    id="1",
-                    content=[
-                        ResponseOutputText(
-                            annotations=[],
-                            text="This is a haiku (not really)",
-                            type="output_text",
-                        )
-                    ],
-                    role="assistant",
-                    status="completed",
-                    type="message",
-                )
-            ],
-            usage=Usage(requests=1, input_tokens=1, output_tokens=1, total_tokens=1),
-            response_id="1",
-        )
-    ]
+def test_model():
+    return sequential_test_model(
+        [
+            ModelResponse(
+                output=[
+                    ResponseOutputMessage(
+                        id="1",
+                        content=[
+                            ResponseOutputText(
+                                annotations=[],
+                                text="This is a haiku (not really)",
+                                type="output_text",
+                            )
+                        ],
+                        role="assistant",
+                        status="completed",
+                        type="message",
+                    )
+                ],
+                usage=Usage(
+                    requests=1, input_tokens=1, output_tokens=1, total_tokens=1
+                ),
+                response_id="1",
+            )
+        ]
+    )
+
 
-    return mock
+@pytest.fixture
+def test_model():
+    return TestModel.returning_responses(
+        [ResponseBuilders.output_message("This is a haiku (not really)")]
+    )
 
 
 async def test_execute_workflow(client: Client):

tests/openai_agents/conftest.py

@@ -2,21 +2,23 @@
 from typing import Optional
 
 import pytest
-from agents import Model, ModelProvider
-from temporalio.contrib.openai_agents import ModelActivityParameters, OpenAIAgentsPlugin
+from agents import ModelProvider, ModelResponse
+from temporalio.contrib.openai_agents import (
+    ModelActivityParameters,
+    OpenAIAgentsPlugin,
+    TestModel,
+    TestModelProvider,
+)
 
 
-class MockedModelProvider(ModelProvider):
-    def __init__(self, mocked_model):
-        self.mocked_model = mocked_model
-
-    def get_model(self, model_name: str | None) -> Model:
-        return self.mocked_model
+def sequential_test_model(responses: list[ModelResponse]) -> TestModel:
+    responses = iter(responses)
+    return TestModel(lambda: next(responses))
 
 
 @pytest.fixture
-def model_provider(mocked_model):
-    return MockedModelProvider(mocked_model)
+def model_provider(test_model):
+    return TestModelProvider(test_model)
 
 
 @pytest.fixture