fix unit test for adding new columns with foreign keys

2 files changed, 20 insertions, 4 deletions

diff --git a/migrate/changeset/ansisql.py b/migrate/changeset/ansisql.py
index e6cdd7c..e151090 100644
--- a/migrate/changeset/ansisql.py
+++ b/migrate/changeset/ansisql.py
@@ -17,7 +17,7 @@ from sqlalchemy.schema import (ForeignKeyConstraint,
                                Index)
 
 from migrate import exceptions
-from migrate.changeset import constraint, SQLA_06
+from migrate.changeset import constraint, SQLA_06, SQLA_07
 
 if not SQLA_06:
     from sqlalchemy.sql.compiler import SchemaGenerator, SchemaDropper
@@ -114,8 +114,17 @@ class ANSIColumnGenerator(AlterTableVisitor, SchemaGenerator):
                                         name=column.unique_name).create()
 
         # SA bounds FK constraints to table, add manually
-        for fk in column.foreign_keys:
-            self.add_foreignkey(fk.constraint)
+        if not SQLA_07:
+            for fk in column.foreign_keys:
+                self.add_foreignkey(fk.constraint)
+        else:
+            for fk in column.foreign_keys:
+                self.traverse_single(fk)
+            #for fk in list(column.foreign_keys):
+            #    fk.constraint = ForeignKeyConstraint([column],
+            #        [fk.target_fullname], table=column.table)
+            #    self.add_foreignkey(fk.constraint)
+            #    import pdb; pdb.set_trace()
 
         # add primary key constraint if needed
         if column.primary_key_name:
diff --git a/migrate/tests/changeset/test_changeset.py b/migrate/tests/changeset/test_changeset.py
index d1b34cd..89225a4 100644
--- a/migrate/tests/changeset/test_changeset.py
+++ b/migrate/tests/changeset/test_changeset.py
@@ -169,6 +169,8 @@ class TestAddDropColumn(fixture.DB):
 
         # create column with fk
         col = Column('data', Integer, ForeignKey(reftable.c.id))
+        if SQLA_07:
+            self.table.append_column(col)
         col.create(self.table)
 
         # check if constraint is added
@@ -180,7 +182,12 @@ class TestAddDropColumn(fixture.DB):
 
         # TODO: test on db level if constraints work
 
-        self.assertEqual(reftable.c.id.name, col.foreign_keys[0].column.name)
+        if SQLA_07:
+            self.assertEqual(reftable.c.id.name,
+                list(col.foreign_keys)[0].column.name)
+        else:
+            self.assertEqual(reftable.c.id.name,
+                col.foreign_keys[0].column.name)
         col.drop(self.table)
 
         if self.engine.has_table(reftable.name):