1+ from datetime import datetime , timezone
12from unittest .mock import Mock , patch
23
34import pytest
@@ -17,6 +18,11 @@ class _TestOutput(BaseModel):
1718 message : str
1819
1920
21+ class _TestOutputWithDatetime (BaseModel ):
22+ timestamp : datetime
23+ message : str
24+
25+
2026@pytest .fixture
2127def mock_agent () -> Mock :
2228 mock = Mock (spec = _AgentBase )
@@ -169,6 +175,29 @@ def test_format_output_tool_call_requests(self):
169175URL: https://workflowai.hello/_/agents/agent-id/runs/run-id"""
170176 )
171177
178+ def test_format_output_with_datetime (self ):
179+ """Test that datetimes in the output model are correctly serialized to ISO strings."""
180+ test_dt = datetime (2024 , 1 , 1 , 12 , 30 , 0 , tzinfo = timezone .utc )
181+ run = Run [_TestOutputWithDatetime ](
182+ id = "run-dt-id" ,
183+ agent_id = "agent-dt-id" ,
184+ schema_id = 2 ,
185+ output = _TestOutputWithDatetime (timestamp = test_dt , message = "datetime test" ),
186+ duration_seconds = 0.5 ,
187+ cost_usd = 0.0001 ,
188+ )
189+
190+ expected_json_part = '{\n "timestamp": "2024-01-01T12:30:00",\n "message": "datetime test"\n }'
191+ expected = f"""\n Output:
192+ ==================================================
193+ { expected_json_part }
194+ ==================================================
195+ Cost: $ 0.00010
196+ Latency: 0.50s
197+ URL: https://workflowai.hello/_/agents/agent-dt-id/runs/run-dt-id"""
198+
199+ assert run .format_output () == expected
200+
172201
173202class TestRunURL :
174203 # The @patch decorator from unittest.mock temporarily replaces the value of an attribute
0 commit comments