fix errors in test_pydantic.py

piccolo-orm · Dec 23, 2023 · b2e6cc0 · b2e6cc0
1 parent 18a87c0
commit b2e6cc0
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/tests/utils/test_pydantic.py b/tests/utils/test_pydantic.py
@@ -646,7 +646,7 @@ class Concert(Table):
         # Test two levels deep
 
         BandModel = create_pydantic_model(
-            table=Band, nested=(Band.manager.country,)
+            table=Band, nested=(Band.manager._.country,)
         )
 
         ManagerModel = t.cast(
@@ -679,7 +679,7 @@ class Concert(Table):
         # Test three levels deep
 
         ConcertModel = create_pydantic_model(
-            Concert, nested=(Concert.band_1.manager,)
+            Concert, nested=(Concert.band_1._.manager,)
         )
 
         VenueModel = ConcertModel.model_fields["venue"].annotation
@@ -720,11 +720,14 @@ class Concert(Table):
 
         MyConcertModel = create_pydantic_model(
             Concert,
-            nested=(Concert.band_1.manager,),
+            nested=(Concert.band_1._.manager,),
             model_name="MyConcertModel",
         )
 
-        BandModel = MyConcertModel.model_fields["band_1"].annotation
+        BandModel = t.cast(
+            t.Type[pydantic.BaseModel],
+            MyConcertModel.model_fields["band_1"].annotation,
+        )
         self.assertEqual(BandModel.__qualname__, "MyConcertModel.band_1")
 
         ManagerModel = BandModel.model_fields["manager"].annotation