upgrade more tests

Author: donald-pinckney

tests/openai_agents/basic/test_agent_lifecycle_workflow.py

@@ -2,31 +2,40 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from temporalio.client import Client
+from temporalio.contrib.openai_agents.testing import AgentEnvironment, ResponseBuilders, TestModel
 from temporalio.worker import Worker
 
 from openai_agents.basic.workflows.agent_lifecycle_workflow import (
     AgentLifecycleWorkflow,
 )
 
 
+def agent_lifecycle_test_model():
+    return TestModel.returning_responses(
+        [ResponseBuilders.output_message('{"number": 10}')]
+    )
+
+
 async def test_execute_workflow(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[AgentLifecycleWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-    ):
-        result = await client.execute_workflow(
-            AgentLifecycleWorkflow.run,
-            10,  # max_number parameter
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=agent_lifecycle_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
+            workflows=[AgentLifecycleWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+        ):
+            result = await client.execute_workflow(
+                AgentLifecycleWorkflow.run,
+                10,  # max_number parameter
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
 
-        # Verify the result has the expected structure
-        assert isinstance(result.number, int)
-        assert (
-            0 <= result.number <= 20
-        )  # Should be between 0 and max*2 due to multiply operation
+            # Verify the result has the expected structure
+            assert isinstance(result.number, int)
+            assert (
+                0 <= result.number <= 20
+            )  # Should be between 0 and max*2 due to multiply operation

tests/openai_agents/basic/test_dynamic_system_prompt_workflow.py

@@ -2,51 +2,62 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from temporalio.client import Client
+from temporalio.contrib.openai_agents.testing import AgentEnvironment, ResponseBuilders, TestModel
 from temporalio.worker import Worker
 
 from openai_agents.basic.workflows.dynamic_system_prompt_workflow import (
     DynamicSystemPromptWorkflow,
 )
 
 
+def dynamic_system_prompt_test_model():
+    return TestModel.returning_responses(
+        [ResponseBuilders.output_message("Style: haiku\nResponse: The weather is cloudy with a chance of meatballs.")]
+    )
+
+
 async def test_execute_workflow_with_random_style(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[DynamicSystemPromptWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-    ):
-        result = await client.execute_workflow(
-            DynamicSystemPromptWorkflow.run,
-            "Tell me about the weather today.",
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=dynamic_system_prompt_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
-
-        # Verify the result has the expected format
-        assert "Style:" in result
-        assert "Response:" in result
-        assert any(style in result for style in ["haiku", "pirate", "robot"])
+            workflows=[DynamicSystemPromptWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+        ):
+            result = await client.execute_workflow(
+                DynamicSystemPromptWorkflow.run,
+                "Tell me about the weather today.",
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
+
+            # Verify the result has the expected format
+            assert "Style:" in result
+            assert "Response:" in result
+            assert any(style in result for style in ["haiku", "pirate", "robot"])
 
 
 async def test_execute_workflow_with_specific_style(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[DynamicSystemPromptWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-    ):
-        result = await client.execute_workflow(
-            DynamicSystemPromptWorkflow.run,
-            args=["Tell me about the weather today.", "haiku"],
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=dynamic_system_prompt_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
-
-        # Verify the result has the expected format and style
-        assert "Style: haiku" in result
-        assert "Response:" in result
+            workflows=[DynamicSystemPromptWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+        ):
+            result = await client.execute_workflow(
+                DynamicSystemPromptWorkflow.run,
+                args=["Tell me about the weather today.", "haiku"],
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
+
+            # Verify the result has the expected format and style
+            assert "Style: haiku" in result
+            assert "Response:" in result

tests/openai_agents/basic/test_lifecycle_workflow.py

@@ -2,29 +2,38 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from temporalio.client import Client
+from temporalio.contrib.openai_agents.testing import AgentEnvironment, ResponseBuilders, TestModel
 from temporalio.worker import Worker
 
 from openai_agents.basic.workflows.lifecycle_workflow import LifecycleWorkflow
 
 
+def lifecycle_test_model():
+    return TestModel.returning_responses(
+        [ResponseBuilders.output_message('{"number": 10}')]
+    )
+
+
 async def test_execute_workflow(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[LifecycleWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-    ):
-        result = await client.execute_workflow(
-            LifecycleWorkflow.run,
-            10,  # max_number parameter
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=lifecycle_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
+            workflows=[LifecycleWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+        ):
+            result = await client.execute_workflow(
+                LifecycleWorkflow.run,
+                10,  # max_number parameter
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
 
-        # Verify the result has the expected structure
-        assert isinstance(result.number, int)
-        assert (
-            0 <= result.number <= 20
-        )  # Should be between 0 and max*2 due to multiply operation
+            # Verify the result has the expected structure
+            assert isinstance(result.number, int)
+            assert (
+                0 <= result.number <= 20
+            )  # Should be between 0 and max*2 due to multiply operation

tests/openai_agents/basic/test_local_image_workflow.py

@@ -2,52 +2,63 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from temporalio.client import Client
+from temporalio.contrib.openai_agents.testing import AgentEnvironment, ResponseBuilders, TestModel
 from temporalio.worker import Worker
 
 from openai_agents.basic.activities.image_activities import read_image_as_base64
 from openai_agents.basic.workflows.local_image_workflow import LocalImageWorkflow
 
 
+def local_image_test_model():
+    return TestModel.returning_responses(
+        [ResponseBuilders.output_message("I can see a bison in the image.")]
+    )
+
+
 async def test_execute_workflow_default_question(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[LocalImageWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-        activities=[read_image_as_base64],
-    ):
-        result = await client.execute_workflow(
-            LocalImageWorkflow.run,
-            "openai_agents/basic/media/image_bison.jpg",  # Path to test image
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=local_image_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
+            workflows=[LocalImageWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+            activities=[read_image_as_base64],
+        ):
+            result = await client.execute_workflow(
+                LocalImageWorkflow.run,
+                "openai_agents/basic/media/image_bison.jpg",  # Path to test image
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
 
-        # Verify the result is a string response
-        assert isinstance(result, str)
-        assert len(result) > 0
+            # Verify the result is a string response
+            assert isinstance(result, str)
+            assert len(result) > 0
 
 
 async def test_execute_workflow_custom_question(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[LocalImageWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-        activities=[read_image_as_base64],
-    ):
-        custom_question = "What animals do you see in this image?"
-        result = await client.execute_workflow(
-            LocalImageWorkflow.run,
-            args=["openai_agents/basic/media/image_bison.jpg", custom_question],
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=local_image_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
+            workflows=[LocalImageWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+            activities=[read_image_as_base64],
+        ):
+            custom_question = "What animals do you see in this image?"
+            result = await client.execute_workflow(
+                LocalImageWorkflow.run,
+                args=["openai_agents/basic/media/image_bison.jpg", custom_question],
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
 
-        # Verify the result is a string response
-        assert isinstance(result, str)
-        assert len(result) > 0
+            # Verify the result is a string response
+            assert isinstance(result, str)
+            assert len(result) > 0

tests/openai_agents/basic/test_non_strict_output_workflow.py

@@ -2,41 +2,51 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from temporalio.client import Client
+from temporalio.contrib.openai_agents.testing import AgentEnvironment, ResponseBuilders, TestModel
 from temporalio.worker import Worker
 
 from openai_agents.basic.workflows.non_strict_output_workflow import (
     NonStrictOutputWorkflow,
 )
 
 
+def non_strict_output_test_model():
+    return TestModel.returning_responses(
+        [ResponseBuilders.output_message('{"jokes": {"1": "Why do programmers prefer dark mode? Because light attracts bugs!", "2": "How many programmers does it take to change a light bulb? None, that\'s a hardware problem.", "3": "Why do Java developers wear glasses? Because they can\'t C#!"}}')
+        ]
+    )
+
+
 async def test_execute_workflow(client: Client):
     task_queue_name = str(uuid.uuid4())
 
-    async with Worker(
-        client,
-        task_queue=task_queue_name,
-        workflows=[NonStrictOutputWorkflow],
-        activity_executor=ThreadPoolExecutor(5),
-        # No external activities needed
-    ):
-        result = await client.execute_workflow(
-            NonStrictOutputWorkflow.run,
-            "Tell me 3 funny jokes about programming.",
-            id=str(uuid.uuid4()),
+    async with AgentEnvironment(model=non_strict_output_test_model()) as agent_env:
+        client = agent_env.applied_on_client(client)
+        async with Worker(
+            client,
             task_queue=task_queue_name,
-        )
-
-        # Verify the result has the expected structure
-        assert isinstance(result, dict)
-
-        assert "strict_error" in result
-        assert "non_strict_result" in result
-
-        # If there's a strict_error, it should be a string
-        if "strict_error" in result:
-            assert isinstance(result["strict_error"], str)
-            assert len(result["strict_error"]) > 0
-
-        jokes = result["non_strict_result"]["jokes"]
-        assert isinstance(jokes, dict)
-        assert isinstance(jokes[list(jokes.keys())[0]], str)
+            workflows=[NonStrictOutputWorkflow],
+            activity_executor=ThreadPoolExecutor(5),
+            # No external activities needed
+        ):
+            result = await client.execute_workflow(
+                NonStrictOutputWorkflow.run,
+                "Tell me 3 funny jokes about programming.",
+                id=str(uuid.uuid4()),
+                task_queue=task_queue_name,
+            )
+
+            # Verify the result has the expected structure
+            assert isinstance(result, dict)
+
+            assert "strict_error" in result
+            assert "non_strict_result" in result
+
+            # If there's a strict_error, it should be a string
+            if "strict_error" in result:
+                assert isinstance(result["strict_error"], str)
+                assert len(result["strict_error"]) > 0
+
+            jokes = result["non_strict_result"]["jokes"]
+            assert isinstance(jokes, dict)
+            assert isinstance(jokes[list(jokes.keys())[0]], str)