fix LazyTableReference with ModelBuilder (#870)

piccolo/testing/model_builder.py

@@ -6,7 +6,7 @@
 from decimal import Decimal
 from uuid import UUID
 
-from piccolo.columns import Array, Column
+from piccolo.columns import Array, Column, ForeignKey
 from piccolo.custom_types import TableInstance
 from piccolo.testing.random_builder import RandomBuilder
 from piccolo.utils.sync import run_sync
@@ -122,9 +122,9 @@ async def _build(
             if column._meta.name in defaults:
                 continue  # Column value exists
 
-            if "references" in column._meta.params and persist:
+            if isinstance(column, ForeignKey) and persist:
                 reference_model = await cls._build(
-                    column._meta.params["references"],
+                    column._foreign_key_meta.resolved_references,
                     persist=True,
                 )
                 random_value = getattr(

tests/testing/test_model_builder.py

@@ -1,8 +1,18 @@
 import asyncio
+import typing as t
 import unittest
 
-from piccolo.columns import Array, Decimal, Integer, Numeric, Real, Varchar
-from piccolo.table import Table
+from piccolo.columns import (
+    Array,
+    Decimal,
+    ForeignKey,
+    Integer,
+    LazyTableReference,
+    Numeric,
+    Real,
+    Varchar,
+)
+from piccolo.table import Table, create_db_tables_sync, drop_db_tables_sync
 from piccolo.testing.model_builder import ModelBuilder
 from tests.base import engines_skip
 from tests.example_apps.music.tables import (
@@ -30,62 +40,52 @@ class TableWithDecimal(Table):
     decimal_with_digits = Decimal(digits=(4, 2))
 
 
+class BandWithLazyReference(Table):
+    manager = ForeignKey(
+        references=LazyTableReference(
+            "Manager", module_path="tests.example_apps.music.tables"
+        )
+    )
+
+
+TABLES = (
+    Manager,
+    Band,
+    Poster,
+    RecordingStudio,
+    Shirt,
+    Venue,
+    Concert,
+    Ticket,
+    TableWithArrayField,
+    TableWithDecimal,
+    BandWithLazyReference,
+)
+
+
 # Cockroach Bug: Can turn ON when resolved: https://github.com/cockroachdb/cockroach/issues/71908  # noqa: E501
 @engines_skip("cockroach")
 class TestModelBuilder(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-
-        for table_class in (
-            Manager,
-            Band,
-            Poster,
-            RecordingStudio,
-            Shirt,
-            Venue,
-            Concert,
-            Ticket,
-            TableWithArrayField,
-            TableWithDecimal,
-        ):
-            table_class.create_table().run_sync()
+        create_db_tables_sync(*TABLES)
 
     @classmethod
     def tearDownClass(cls) -> None:
-        for table_class in (
-            TableWithDecimal,
-            TableWithArrayField,
-            Ticket,
-            Concert,
-            Venue,
-            Shirt,
-            RecordingStudio,
-            Poster,
-            Band,
-            Manager,
-        ):
-            table_class.alter().drop_table().run_sync()
-
-    def test_model_builder_async(self):
-        async def build_model(model):
-            return await ModelBuilder.build(model)
-
-        asyncio.run(build_model(Manager))
-        asyncio.run(build_model(Ticket))
-        asyncio.run(build_model(Poster))
-        asyncio.run(build_model(RecordingStudio))
-        asyncio.run(build_model(TableWithArrayField))
-        asyncio.run(build_model(TableWithDecimal))
-
-    def test_model_builder_sync(self):
-        ModelBuilder.build_sync(Manager)
-        ModelBuilder.build_sync(Ticket)
-        ModelBuilder.build_sync(Poster)
-        ModelBuilder.build_sync(RecordingStudio)
-        ModelBuilder.build_sync(TableWithArrayField)
-        ModelBuilder.build_sync(TableWithDecimal)
-
-    def test_model_builder_with_choices(self):
+        drop_db_tables_sync(*TABLES)
+
+    def test_async(self):
+        async def build_model(table_class: t.Type[Table]):
+            return await ModelBuilder.build(table_class)
+
+        for table_class in TABLES:
+            asyncio.run(build_model(table_class))
+
+    def test_sync(self):
+        for table_class in TABLES:
+            ModelBuilder.build_sync(table_class)
+
+    def test_choices(self):
         shirt = ModelBuilder.build_sync(Shirt)
         queried_shirt = (
             Shirt.objects().where(Shirt.id == shirt.id).first().run_sync()
@@ -96,30 +96,35 @@ def test_model_builder_with_choices(self):
             ["s", "l", "m"],
         )
 
-    def test_model_builder_with_foreign_key(self):
-        ModelBuilder.build_sync(Band)
+    def test_foreign_key(self):
+        model = ModelBuilder.build_sync(Band, persist=True)
 
-    def test_model_builder_with_invalid_column(self):
+        self.assertTrue(
+            Manager.exists().where(Manager.id == model.manager).run_sync()
+        )
+
+    def test_lazy_foreign_key(self):
+        model = ModelBuilder.build_sync(BandWithLazyReference, persist=True)
+
+        self.assertTrue(
+            Manager.exists().where(Manager.id == model.manager).run_sync()
+        )
+
+    def test_invalid_column(self):
         with self.assertRaises(ValueError):
             ModelBuilder.build_sync(Band, defaults={"X": 1})
 
-    def test_model_builder_with_minimal(self):
+    def test_minimal(self):
         band = ModelBuilder.build_sync(Band, minimal=True)
 
-        self.assertEqual(
-            Band.exists().where(Band.id == band.id).run_sync(),
-            True,
-        )
+        self.assertTrue(Band.exists().where(Band.id == band.id).run_sync())
 
-    def test_model_builder_with_no_persist(self):
+    def test_persist_false(self):
         band = ModelBuilder.build_sync(Band, persist=False)
 
-        self.assertEqual(
-            Band.exists().where(Band.id == band.id).run_sync(),
-            False,
-        )
+        self.assertFalse(Band.exists().where(Band.id == band.id).run_sync())
 
-    def test_model_builder_with_valid_column(self):
+    def test_valid_column(self):
         manager = ModelBuilder.build_sync(
             Manager, defaults={Manager.name: "Guido"}
         )
@@ -133,7 +138,7 @@ def test_model_builder_with_valid_column(self):
 
         self.assertEqual(queried_manager.name, "Guido")
 
-    def test_model_builder_with_valid_column_string(self):
+    def test_valid_column_string(self):
         manager = ModelBuilder.build_sync(Manager, defaults={"name": "Guido"})
 
         queried_manager = (
@@ -145,14 +150,14 @@ def test_model_builder_with_valid_column_string(self):
 
         self.assertEqual(queried_manager.name, "Guido")
 
-    def test_model_builder_with_valid_foreign_key(self):
+    def test_valid_foreign_key(self):
         manager = ModelBuilder.build_sync(Manager)
 
         band = ModelBuilder.build_sync(Band, defaults={Band.manager: manager})
 
         self.assertEqual(manager._meta.primary_key, band.manager)
 
-    def test_model_builder_with_valid_foreign_key_string(self):
+    def test_valid_foreign_key_string(self):
         manager = ModelBuilder.build_sync(Manager)
 
         band = ModelBuilder.build_sync(Band, defaults={"manager": manager})