@@ -195,3 +195,158 @@ def test_upgrade_scenario_basic_to_current(self) -> None:
195195 # Missing tables
196196 assert "password_reset_token" in added_table_names
197197 assert "email_verification_token" in added_table_names
198+
199+
200+ class TestMigrateFixMigrationRendering :
201+ """Test that add_table diffs render correct migration content."""
202+
203+ def _render_add_table_upgrade (self , table : sa .Table ) -> tuple [list [str ], list [str ]]:
204+ """Simulate the migration rendering logic from migrate_fix.py."""
205+
206+ def _sa_type_str (col_type : object ) -> str :
207+ type_name = type (col_type ).__name__
208+ return f"sa.{ type_name }
209+
210+ upgrade_lines : list [str ] = []
211+ downgrade_lines : list [str ] = []
212+
213+ cols = []
214+ for col in table .columns :
215+ col_str = f"sa.Column('{ col .name } { _sa_type_str (col .type )}
216+ col_str += f", nullable={ col .nullable }
217+ if col .primary_key :
218+ col_str += ", primary_key=True"
219+ if col .server_default is not None :
220+ col_str += f", server_default=sa.text('{ col .server_default .arg } # type: ignore[union-attr]
221+ col_str += ")"
222+ cols .append (f" { col_str } )
223+ cols_block = "\n " .join (cols )
224+ upgrade_lines .append (
225+ f" op.create_table(\n '{ table .name } \n { cols_block } \n )"
226+ )
227+
228+ for fk in table .foreign_key_constraints :
229+ local_cols = [c .name for c in fk .columns ]
230+ ref_table = fk .referred_table .name
231+ ref_cols = [c .name for c in fk .referred_table .primary_key .columns ]
232+ fk_name = fk .name or f"fk_{ table .name } { '_' .join (local_cols )} { ref_table }
233+ upgrade_lines .append (
234+ f" op.create_foreign_key('{ fk_name } { table .name } { ref_table } { local_cols } { ref_cols }
235+ )
236+ downgrade_lines .append (
237+ f" op.drop_constraint('{ fk_name } { table .name }
238+ )
239+
240+ for idx in table .indexes :
241+ idx_cols = [c .name for c in idx .columns ]
242+ unique_str = ", unique=True" if idx .unique else ""
243+ upgrade_lines .append (
244+ f" op.create_index('{ idx .name } { table .name } { idx_cols } { unique_str }
245+ )
246+
247+ downgrade_lines .append (f" op.drop_table('{ table .name } )
248+
249+ return upgrade_lines , downgrade_lines
250+
251+ def test_renders_create_table (self ) -> None :
252+ """add_table diff renders op.create_table with all columns."""
253+ metadata = sa .MetaData ()
254+ table = sa .Table (
255+ "password_reset_token" ,
256+ metadata ,
257+ sa .Column ("id" , sa .Integer (), primary_key = True ),
258+ sa .Column ("user_id" , sa .Integer (), nullable = False ),
259+ sa .Column ("token" , sa .String (), nullable = False ),
260+ sa .Column (
261+ "used" , sa .Boolean (), nullable = False , server_default = sa .text ("0" )
262+ ),
263+ )
264+
265+ upgrade , downgrade = self ._render_add_table_upgrade (table )
266+
267+ # Should have create_table
268+ create_line = upgrade [0 ]
269+ assert "op.create_table(" in create_line
270+ assert "'password_reset_token'" in create_line
271+ assert "'id'" in create_line
272+ assert "'user_id'" in create_line
273+ assert "'token'" in create_line
274+ assert "'used'" in create_line
275+ assert "primary_key=True" in create_line
276+ assert "server_default=sa.text('0')" in create_line
277+
278+ # Downgrade drops the table
279+ assert "op.drop_table('password_reset_token')" in downgrade [- 1 ]
280+
281+ def test_renders_foreign_keys (self ) -> None :
282+ """add_table diff with FKs renders op.create_foreign_key."""
283+ metadata = sa .MetaData ()
284+ sa .Table (
285+ "user" ,
286+ metadata ,
287+ sa .Column ("id" , sa .Integer (), primary_key = True ),
288+ )
289+ sa .Table (
290+ "password_reset_token" ,
291+ metadata ,
292+ sa .Column ("id" , sa .Integer (), primary_key = True ),
293+ sa .Column (
294+ "user_id" ,
295+ sa .Integer (),
296+ sa .ForeignKey ("user.id" ),
297+ nullable = False ,
298+ ),
299+ sa .Column ("token" , sa .String (), nullable = False ),
300+ )
301+
302+ token_table = metadata .tables ["password_reset_token" ]
303+ upgrade , downgrade = self ._render_add_table_upgrade (token_table )
304+
305+ # Should have create_foreign_key
306+ fk_lines = [line for line in upgrade if "create_foreign_key" in line ]
307+ assert len (fk_lines ) == 1
308+ assert "'password_reset_token'" in fk_lines [0 ]
309+ assert "'user'" in fk_lines [0 ]
310+ assert "['user_id']" in fk_lines [0 ]
311+ assert "['id']" in fk_lines [0 ]
312+
313+ # Downgrade should drop the constraint
314+ drop_fk_lines = [line for line in downgrade if "drop_constraint" in line ]
315+ assert len (drop_fk_lines ) == 1
316+ assert "type_='foreignkey'" in drop_fk_lines [0 ]
317+
318+ def test_renders_indexes (self ) -> None :
319+ """add_table diff with indexes renders op.create_index."""
320+ metadata = sa .MetaData ()
321+ table = sa .Table (
322+ "org_invite" ,
323+ metadata ,
324+ sa .Column ("id" , sa .Integer (), primary_key = True ),
325+ sa .Column ("email" , sa .String (), nullable = False ),
326+ sa .Column ("token" , sa .String (), nullable = False ),
327+ )
328+ sa .Index ("ix_org_invite_token" , table .c .token , unique = True )
329+
330+ upgrade , _ = self ._render_add_table_upgrade (table )
331+
332+ idx_lines = [line for line in upgrade if "create_index" in line ]
333+ assert len (idx_lines ) == 1
334+ assert "'ix_org_invite_token'" in idx_lines [0 ]
335+ assert "'org_invite'" in idx_lines [0 ]
336+ assert "unique=True" in idx_lines [0 ]
337+
338+ def test_empty_table_renders (self ) -> None :
339+ """Table with only a PK column still renders correctly."""
340+ metadata = sa .MetaData ()
341+ table = sa .Table (
342+ "simple" ,
343+ metadata ,
344+ sa .Column ("id" , sa .Integer (), primary_key = True ),
345+ )
346+
347+ upgrade , downgrade = self ._render_add_table_upgrade (table )
348+
349+ assert "op.create_table(" in upgrade [0 ]
350+ assert "'simple'" in upgrade [0 ]
351+ assert len (upgrade ) == 1 # No FKs or indexes
352+ assert "op.drop_table('simple')" in downgrade [- 1 ]
0 commit comments