From 2a71180c1d1b578fa9777043fdcb30549e522ff6 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Mon, 17 Mar 2025 12:43:21 -0700
Subject: [PATCH] Add tests for idempotency in migraionts

---
 libs/checkpoint-postgres/tests/test_async_store.py | 4 ++++
 libs/checkpoint-postgres/tests/test_store.py       | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py
index 09502403d..94dc1fa03 100644
--- a/libs/checkpoint-postgres/tests/test_async_store.py
+++ b/libs/checkpoint-postgres/tests/test_async_store.py
@@ -69,6 +69,10 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
                 for mig in store.MIGRATIONS
             ]
             await store.setup()
+            async with store._cursor() as cur:
+                # drop the migration index
+                await cur.execute("DROP TABLE IF EXISTS store_migrations")
+            await store.setup()  # Will fail if migrations aren't idempotent
 
         if request.param == "pipe":
             async with AsyncPostgresStore.from_conn_string(
diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py
index 4ee37484e..a1360174a 100644
--- a/libs/checkpoint-postgres/tests/test_store.py
+++ b/libs/checkpoint-postgres/tests/test_store.py
@@ -415,6 +415,10 @@ def _create_vector_store(
             ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
         ) as store:
             store.setup()
+            with store._cursor() as cur:
+                # drop the migration index
+                cur.execute("DROP TABLE IF EXISTS store_migrations")
+            store.setup()  # Will fail if migrations aren't idempotent
             yield store
     finally:
         with Connection.connect(admin_conn_string, autocommit=True) as conn: