1
+ from datetime import datetime , timezone
1
2
from unittest .mock import Mock , patch
2
3
3
4
import pytest
@@ -17,6 +18,11 @@ class _TestOutput(BaseModel):
17
18
message : str
18
19
19
20
21
+ class _TestOutputWithDatetime (BaseModel ):
22
+ timestamp : datetime
23
+ message : str
24
+
25
+
20
26
@pytest .fixture
21
27
def mock_agent () -> Mock :
22
28
mock = Mock (spec = _AgentBase )
@@ -169,6 +175,29 @@ def test_format_output_tool_call_requests(self):
169
175
URL: https://workflowai.hello/_/agents/agent-id/runs/run-id"""
170
176
)
171
177
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:00Z",\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
+
172
201
173
202
class TestRunURL :
174
203
# The @patch decorator from unittest.mock temporarily replaces the value of an attribute
0 commit comments