Fix crutch in CompareAndRunTestCase

Pogchamp-company · Mar 17, 2024 · 34aee50 · 34aee50
1 parent a1d58ff
commit 34aee50
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/tests/base/run_migration_test_abc.py b/tests/base/run_migration_test_abc.py
@@ -17,7 +17,7 @@ def get_database_schema(self) -> MetaData:
     def get_target_schema(self) -> MetaData:
         ...
 
-    def insert_migration_data(self, connection: "Connection"):
+    def insert_migration_data(self, connection: "Connection", database_schema: MetaData) -> None:
         pass
 
     @abstractmethod
@@ -33,7 +33,7 @@ def test_run(self, connection: "Connection"):
         target_schema = self.get_target_schema()
 
         database_schema.create_all(connection)
-        self.insert_migration_data(connection)
+        self.insert_migration_data(connection, database_schema)
 
         compare_and_run(
             connection,

diff --git a/tests/test_alter_column/test_text_column.py b/tests/test_alter_column/test_text_column.py
@@ -19,17 +19,18 @@ class NewEnum(Enum):
 class TestTextColumn(CompareAndRunTestCase):
     def get_database_schema(self) -> MetaData:
         database_schema = MetaData()
-        self.a_table = Table("a", database_schema, Column("value", TEXT))
+        Table("a", database_schema, Column("value", TEXT))
         return database_schema
 
     def get_target_schema(self) -> MetaData:
         target_schema = MetaData()
         Table("a", target_schema, Column("value", postgresql.ENUM(NewEnum)))
         return target_schema
 
-    def insert_migration_data(self, connection: "Connection"):
+    def insert_migration_data(self, connection: "Connection", database_schema: MetaData) -> None:
+        a_table = database_schema.tables["a"]
         connection.execute(
-            insert(self.a_table).values(
+            insert(a_table).values(
                 [
                     {"value": NewEnum.A.name},
                     {"value": NewEnum.B.name},