diff --git a/tests/test_migrate_legacy_data.py b/tests/test_migrate_legacy_data.py
index 88dedd5..ca24b53 100644
--- a/tests/test_migrate_legacy_data.py
+++ b/tests/test_migrate_legacy_data.py
@@ -277,7 +277,7 @@ def test_migrate_reconciliation_failure_raises(tmp_path: Path, monkeypatch: pyte
     _make_legacy_location_db(legacy_path, LOCATION_ROWS)
     legacy_url = f"sqlite:///{legacy_path}"
 
-    def _always_fail(conn, table, pk_cols, source_count):
+    def _always_fail(conn, table, columns, source_count):
         # Simulate a scenario where reconciliation finds rows missing
         raise RuntimeError(
             f"Reconciliation failed for table '{table}': "
@@ -302,7 +302,7 @@ def test_cli_exits_nonzero_on_reconciliation_failure(
     legacy_url = f"sqlite:///{legacy_path}"
 
     # Patch _reconcile to always raise
-    def _always_fail(conn, table, pk_cols, source_count):
+    def _always_fail(conn, table, columns, source_count):
         raise RuntimeError(
             f"Reconciliation failed for table '{table}': 1 row missing."
         )