From 53a8b33cb72ecf68a2ab918a13c1b2ea6c3adb42 Mon Sep 17 00:00:00 2001
From: Ben Cail <bcail@crossway.org>
Date: Mon, 4 Nov 2024 13:23:16 -0500
Subject: [PATCH] updates from previous PR changes

---
 tests/migrations/test_base.py       |  2 +-
 tests/migrations/test_operations.py | 17 ++++++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/tests/migrations/test_base.py b/tests/migrations/test_base.py
index cde4063d04..bb9f2bb234 100644
--- a/tests/migrations/test_base.py
+++ b/tests/migrations/test_base.py
@@ -104,7 +104,7 @@ class MigrationTestBase(TransactionTestCase):
                     .values()
                     if (
                         c["columns"] == list(columns)
-                        and (index_type is None or c["type"] == index_type)
+                        and (index_type is None or c.get("type") == index_type)
                         and not c["unique"]
                     )
                 ),
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 318cc527c2..8d58f8b2ad 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -6235,11 +6235,12 @@ class PrimaryKeyOperations(OperationTestBase):
             "field2",
             models.SlugField(max_length=20, primary_key=True),
         )
-
         project_state = ProjectState()
+        new_state = project_state.clone()
+        operation1.state_forwards("migrtest", new_state)
+
+        self.assertTableNotExists("migrtest_simplemodel")
         with connection.schema_editor() as editor:
-            new_state = project_state.clone()
-            operation1.state_forwards("migrtest", new_state)
             operation1.database_forwards("migrtest", editor, project_state, new_state)
             project_state, new_state = new_state, new_state.clone()
             operation2.state_forwards("migrtest", new_state)
@@ -6247,6 +6248,16 @@ class PrimaryKeyOperations(OperationTestBase):
             project_state, new_state = new_state, new_state.clone()
             operation3.state_forwards("migrtest", new_state)
             operation3.database_forwards("migrtest", editor, project_state, new_state)
+        self.assertTableExists("migrtest_simplemodel")
+        self.assertColumnExists("migrtest_simplemodel", "field1")
+        self.assertColumnExists("migrtest_simplemodel", "field2")
+        with connection.cursor() as cursor:
+            primary_keys = connection.introspection.get_primary_key_columns(
+                cursor, "migrtest_simplemodel"
+            )
+        self.assertEqual(["field2"], primary_keys)
+        self.assertIndexExists("migrtest_simplemodel", ["field1"], index_type="idx")
+        self.assertIndexExists("migrtest_simplemodel", ["field2"], index_type="idx")
 
 
 class SwappableOperationTests(OperationTestBase):