@@ -154,7 +154,41 @@ class FailedCallback(pydantic.BaseModel):
154
154
DogCallback = Annotated [Union [SuccessfulCallback , FailedCallback ], pydantic .Field (discriminator = "status" )]
155
155
156
156
@event_parser (model = DogCallback )
157
- def handler (event : test_input , _ : Any ) -> str :
157
+ def handler (event , _ : Any ) -> str :
158
+ if isinstance (event , FailedCallback ):
159
+ return f"Uh oh. Had a problem: { event .error } "
160
+
161
+ return f"Successfully retrieved { event .breed } named { event .name } "
162
+
163
+ ret = handler (test_input , None )
164
+ assert ret == expected
165
+
166
+
167
+ @pytest .mark .parametrize (
168
+ "test_input,expected" ,
169
+ [
170
+ (
171
+ {"status" : "succeeded" , "name" : "Clifford" , "breed" : "Labrador" },
172
+ "Successfully retrieved Labrador named Clifford" ,
173
+ ),
174
+ ({"status" : "failed" , "error" : "oh some error" }, "Uh oh. Had a problem: oh some error" ),
175
+ ],
176
+ )
177
+ def test_parser_unions_with_type_adapter_instance (test_input , expected ):
178
+ class SuccessfulCallback (pydantic .BaseModel ):
179
+ status : Literal ["succeeded" ]
180
+ name : str
181
+ breed : Literal ["Newfoundland" , "Labrador" ]
182
+
183
+ class FailedCallback (pydantic .BaseModel ):
184
+ status : Literal ["failed" ]
185
+ error : str
186
+
187
+ DogCallback = Annotated [Union [SuccessfulCallback , FailedCallback ], pydantic .Field (discriminator = "status" )]
188
+ DogCallbackTypeAdapter = pydantic .TypeAdapter (DogCallback )
189
+
190
+ @event_parser (model = DogCallbackTypeAdapter )
191
+ def handler (event , _ : Any ) -> str :
158
192
if isinstance (event , FailedCallback ):
159
193
return f"Uh oh. Had a problem: { event .error } "
160
194
0 commit comments