1+ """Test JobResult serialization with date fields."""
2+
3+ import json
4+ from datetime import date
5+ from pydantic import BaseModel , Field
6+
7+ from batchata .core .job_result import JobResult
8+
9+
10+ class MockPropertyData (BaseModel ):
11+ """Mock Pydantic model with date fields for testing."""
12+ property_name : str = Field (..., description = "Property name" )
13+ appraisal_date : date = Field (..., description = "Date appraisal was performed" )
14+ last_sale_date : date | None = Field (None , description = "Date of last sale" )
15+ appraised_value : float | None = Field (None , description = "Appraised value" )
16+
17+
18+ def test_job_result_serialization_with_dates ():
19+ """Test that JobResult can serialize Pydantic models with date fields to JSON."""
20+ # Create a mock Pydantic model with date fields
21+ property_data = MockPropertyData (
22+ property_name = "Test Property" ,
23+ appraisal_date = date (2023 , 10 , 20 ),
24+ last_sale_date = date (2022 , 5 , 15 ),
25+ appraised_value = 1500000.0
26+ )
27+
28+ # Create JobResult with the mock data
29+ job_result = JobResult (
30+ job_id = "test-job-123" ,
31+ raw_response = "Test raw response" ,
32+ parsed_response = property_data ,
33+ input_tokens = 100 ,
34+ output_tokens = 50 ,
35+ cost_usd = 0.15
36+ )
37+
38+ # Test to_dict() method
39+ result_dict = job_result .to_dict ()
40+
41+ # Verify the parsed_response contains serialized dates as strings
42+ assert result_dict ["parsed_response" ]["appraisal_date" ] == "2023-10-20"
43+ assert result_dict ["parsed_response" ]["last_sale_date" ] == "2022-05-15"
44+ assert result_dict ["parsed_response" ]["property_name" ] == "Test Property"
45+ assert result_dict ["parsed_response" ]["appraised_value" ] == 1500000.0
46+
47+ # Test that the result can be serialized to JSON (this was failing before the fix)
48+ json_str = json .dumps (result_dict , indent = 2 )
49+
50+ # Verify we can parse it back
51+ parsed_back = json .loads (json_str )
52+ assert parsed_back ["parsed_response" ]["appraisal_date" ] == "2023-10-20"
53+ assert parsed_back ["parsed_response" ]["last_sale_date" ] == "2022-05-15"
54+
55+ # Test save_to_json method
56+ import tempfile
57+ import os
58+
59+ with tempfile .NamedTemporaryFile (mode = 'w' , suffix = '.json' , delete = False ) as tmp_file :
60+ tmp_path = tmp_file .name
61+
62+ try :
63+ # This should not raise an exception
64+ job_result .save_to_json (tmp_path )
65+
66+ # Verify the file was created and contains valid JSON
67+ with open (tmp_path , 'r' ) as f :
68+ saved_data = json .load (f )
69+
70+ assert saved_data ["parsed_response" ]["appraisal_date" ] == "2023-10-20"
71+ assert saved_data ["parsed_response" ]["last_sale_date" ] == "2022-05-15"
72+
73+ finally :
74+ # Clean up
75+ if os .path .exists (tmp_path ):
76+ os .unlink (tmp_path )
77+
78+
79+ def test_job_result_serialization_with_dict_response ():
80+ """Test that JobResult handles dict responses correctly (should not change)."""
81+ # Test with dict response (no Pydantic model)
82+ dict_response = {
83+ "property_name" : "Test Property" ,
84+ "appraisal_date" : "2023-10-20" , # Already a string
85+ "appraised_value" : 1500000.0
86+ }
87+
88+ job_result = JobResult (
89+ job_id = "test-job-dict" ,
90+ raw_response = "Test raw response" ,
91+ parsed_response = dict_response ,
92+ input_tokens = 100 ,
93+ output_tokens = 50 ,
94+ cost_usd = 0.15
95+ )
96+
97+ # Test to_dict() method
98+ result_dict = job_result .to_dict ()
99+
100+ # Should be unchanged since it's already a dict
101+ assert result_dict ["parsed_response" ]["appraisal_date" ] == "2023-10-20"
102+ assert result_dict ["parsed_response" ]["property_name" ] == "Test Property"
103+
104+ # Should serialize to JSON without issues
105+ json_str = json .dumps (result_dict , indent = 2 )
106+ parsed_back = json .loads (json_str )
107+ assert parsed_back ["parsed_response" ]["appraisal_date" ] == "2023-10-20"
108+
109+
110+ if __name__ == "__main__" :
111+ test_job_result_serialization_with_dates ()
112+ test_job_result_serialization_with_dict_response ()
113+ print ("All tests passed!" )
0 commit comments