Retry delete collection

libs/astradb/pyproject.toml

@@ -82,7 +82,7 @@ build-backend = "poetry.core.masonry.api"
 #
 # https://github.com/tophat/syrupy
 # --snapshot-warn-unused    Prints a warning on unused snapshots rather than fail the test suite.
-addopts = "--snapshot-warn-unused --strict-markers --strict-config --durations=5"
+addopts = "--snapshot-warn-unused --strict-markers --strict-config --durations=5 -vv"
 # Registering custom markers.
 # https://docs.pytest.org/en/7.1.x/example/markers.html#registering-markers
 markers = [

libs/astradb/tests/integration_tests/test_caches.py

@@ -25,6 +25,7 @@
 
 from langchain_astradb import AstraDBCache, AstraDBSemanticCache
 from langchain_astradb.utils.astradb import SetupMode
+from tests.integration_tests.util import retry, aretry
 
 
 class FakeEmbeddings(Embeddings):
@@ -121,7 +122,7 @@ def astradb_cache() -> Iterator[AstraDBCache]:
         namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
     )
     yield cache
-    cache.collection.astra_db.delete_collection("lc_integration_test_cache")
+    retry(cache.collection.astra_db.delete_collection, args=["lc_integration_test_cache"])
 
 
 @pytest.fixture
@@ -134,9 +135,7 @@ async def async_astradb_cache() -> AsyncIterator[AstraDBCache]:
         setup_mode=SetupMode.ASYNC,
     )
     yield cache
-    await cache.async_collection.astra_db.delete_collection(
-        "lc_integration_test_cache_async"
-    )
+    await aretry(cache.async_collection.astra_db.delete_collection, args=["lc_integration_test_cache_async"])
 
 
 @pytest.fixture(scope="module")
@@ -150,7 +149,8 @@ def astradb_semantic_cache() -> Iterator[AstraDBSemanticCache]:
         embedding=fake_embe,
     )
     yield sem_cache
-    sem_cache.collection.astra_db.delete_collection("lc_integration_test_sem_cache")
+    retry(sem_cache.collection.astra_db.delete_collection,
+          args=["lc_integration_test_sem_cache"])
 
 
 @pytest.fixture
@@ -165,13 +165,13 @@ async def async_astradb_semantic_cache() -> AsyncIterator[AstraDBSemanticCache]:
         setup_mode=SetupMode.ASYNC,
     )
     yield sem_cache
-    sem_cache.collection.astra_db.delete_collection(
-        "lc_integration_test_sem_cache_async"
-    )
+    await aretry(sem_cache.async_collection.astra_db.delete_collection,
+                 args=["lc_integration_test_sem_cache_async"])
 
 
 @pytest.mark.skipif(not _has_env_vars(), reason="Missing Astra DB env. vars")
 class TestAstraDBCaches:
+
     def test_astradb_cache(self, astradb_cache: AstraDBCache) -> None:
         self.do_cache_test(FakeLLM(), astradb_cache, "foo")
 

libs/astradb/tests/integration_tests/test_chat_message_histories.py

@@ -9,6 +9,7 @@
     AstraDBChatMessageHistory,
 )
 from langchain_astradb.utils.astradb import SetupMode
+from tests.integration_tests.util import retry, aretry
 
 
 def _has_env_vars() -> bool:
@@ -30,7 +31,7 @@ def history1() -> Iterable[AstraDBChatMessageHistory]:
         namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
     )
     yield history1
-    history1.collection.astra_db.delete_collection("langchain_cmh_test")
+    retry(history1.collection.astra_db.delete_collection, args=["langchain_cmh_test"])
 
 
 @pytest.fixture(scope="function")
@@ -43,7 +44,7 @@ def history2() -> Iterable[AstraDBChatMessageHistory]:
         namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
     )
     yield history2
-    history2.collection.astra_db.delete_collection("langchain_cmh_test")
+    retry(history2.collection.astra_db.delete_collection, args=["langchain_cmh_test"])
 
 
 @pytest.fixture
@@ -57,7 +58,7 @@ async def async_history1() -> AsyncIterable[AstraDBChatMessageHistory]:
         setup_mode=SetupMode.ASYNC,
     )
     yield history1
-    await history1.async_collection.astra_db.delete_collection("langchain_cmh_test")
+    await aretry(history1.async_collection.astra_db.delete_collection, args=["langchain_cmh_test"])
 
 
 @pytest.fixture(scope="function")
@@ -71,7 +72,7 @@ async def async_history2() -> AsyncIterable[AstraDBChatMessageHistory]:
         setup_mode=SetupMode.ASYNC,
     )
     yield history2
-    await history2.async_collection.astra_db.delete_collection("langchain_cmh_test")
+    await aretry(history2.async_collection.astra_db.delete_collection, args=["langchain_cmh_test"])
 
 
 @pytest.mark.skipif(not _has_env_vars(), reason="Missing Astra DB env. vars")

libs/astradb/tests/integration_tests/test_document_loaders.py

@@ -21,6 +21,7 @@
 from astrapy.db import AstraDBCollection, AsyncAstraDBCollection
 
 from langchain_astradb import AstraDBLoader
+from tests.integration_tests.util import retry, aretry
 
 ASTRA_DB_APPLICATION_TOKEN = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
 ASTRA_DB_API_ENDPOINT = os.getenv("ASTRA_DB_API_ENDPOINT")
@@ -49,7 +50,7 @@ def astra_db_collection() -> Iterator[AstraDBCollection]:
 
     yield collection
 
-    astra_db.delete_collection(collection_name)
+    retry(astra_db.delete_collection, args=[collection_name])
 
 
 @pytest.fixture
@@ -70,7 +71,7 @@ async def async_astra_db_collection() -> AsyncIterator[AsyncAstraDBCollection]:
 
     yield collection
 
-    await astra_db.delete_collection(collection_name)
+    await aretry(astra_db.delete_collection, args=[collection_name])
 
 
 @pytest.mark.requires("astrapy")

libs/astradb/tests/integration_tests/test_storage.py

@@ -8,6 +8,7 @@
 
 from langchain_astradb.storage import AstraDBByteStore, AstraDBStore
 from langchain_astradb.utils.astradb import SetupMode
+from tests.integration_tests.util import retry, aretry
 
 
 def _has_env_vars() -> bool:
@@ -74,7 +75,7 @@ def test_mget(self, astra_db: AstraDB) -> None:
             store = init_store(astra_db, collection_name)
             assert store.mget(["key1", "key2"]) == [[0.1, 0.2], "value2"]
         finally:
-            astra_db.delete_collection(collection_name)
+            retry(astra_db.delete_collection, args=[collection_name])
 
     async def test_amget(self, async_astra_db: AsyncAstraDB) -> None:
         """Test AstraDBStore amget method."""
@@ -83,7 +84,7 @@ async def test_amget(self, async_astra_db: AsyncAstraDB) -> None:
             store = await init_async_store(async_astra_db, collection_name)
             assert await store.amget(["key1", "key2"]) == [[0.1, 0.2], "value2"]
         finally:
-            await async_astra_db.delete_collection(collection_name)
+            await aretry(async_astra_db.delete_collection, args=[collection_name])
 
     def test_mset(self, astra_db: AstraDB) -> None:
         """Test that multiple keys can be set with AstraDBStore."""
@@ -95,7 +96,7 @@ def test_mset(self, astra_db: AstraDB) -> None:
             result = store.collection.find_one({"_id": "key2"})
             assert result["data"]["document"]["value"] == "value2"
         finally:
-            astra_db.delete_collection(collection_name)
+            retry(astra_db.delete_collection, args=[collection_name])
 
     async def test_amset(self, async_astra_db: AsyncAstraDB) -> None:
         """Test that multiple keys can be set with AstraDBStore."""
@@ -107,7 +108,7 @@ async def test_amset(self, async_astra_db: AsyncAstraDB) -> None:
             result = await store.async_collection.find_one({"_id": "key2"})
             assert result["data"]["document"]["value"] == "value2"
         finally:
-            await async_astra_db.delete_collection(collection_name)
+            await aretry(async_astra_db.delete_collection, args=[collection_name])
 
     def test_mdelete(self, astra_db: AstraDB) -> None:
         """Test that deletion works as expected."""
@@ -118,7 +119,7 @@ def test_mdelete(self, astra_db: AstraDB) -> None:
             result = store.mget(["key1", "key2"])
             assert result == [None, None]
         finally:
-            astra_db.delete_collection(collection_name)
+            retry(astra_db.delete_collection, args=[collection_name])
 
     async def test_amdelete(self, async_astra_db: AsyncAstraDB) -> None:
         """Test that deletion works as expected."""
@@ -129,7 +130,7 @@ async def test_amdelete(self, async_astra_db: AsyncAstraDB) -> None:
             result = await store.amget(["key1", "key2"])
             assert result == [None, None]
         finally:
-            await async_astra_db.delete_collection(collection_name)
+            await aretry(async_astra_db.delete_collection, args=[collection_name])
 
     def test_yield_keys(self, astra_db: AstraDB) -> None:
         collection_name = "lc_test_store_yield_keys"
@@ -139,7 +140,7 @@ def test_yield_keys(self, astra_db: AstraDB) -> None:
             assert set(store.yield_keys(prefix="key")) == {"key1", "key2"}
             assert set(store.yield_keys(prefix="lang")) == set()
         finally:
-            astra_db.delete_collection(collection_name)
+            retry(astra_db.delete_collection, args=[collection_name])
 
     async def test_ayield_keys(self, async_astra_db: AsyncAstraDB) -> None:
         collection_name = "lc_test_store_yield_keys"
@@ -152,7 +153,7 @@ async def test_ayield_keys(self, async_astra_db: AsyncAstraDB) -> None:
             }
             assert {key async for key in store.ayield_keys(prefix="lang")} == set()
         finally:
-            await async_astra_db.delete_collection(collection_name)
+            await aretry(async_astra_db.delete_collection, args=[collection_name])
 
     def test_bytestore_mget(self, astra_db: AstraDB) -> None:
         """Test AstraDBByteStore mget method."""
@@ -161,7 +162,7 @@ def test_bytestore_mget(self, astra_db: AstraDB) -> None:
             store = init_bytestore(astra_db, collection_name)
             assert store.mget(["key1", "key2"]) == [b"value1", b"value2"]
         finally:
-            astra_db.delete_collection(collection_name)
+            retry(astra_db.delete_collection, args=[collection_name])
 
     def test_bytestore_mset(self, astra_db: AstraDB) -> None:
         """Test that multiple keys can be set with AstraDBByteStore."""
@@ -173,7 +174,7 @@ def test_bytestore_mset(self, astra_db: AstraDB) -> None:
             result = store.collection.find_one({"_id": "key2"})
             assert result["data"]["document"]["value"] == "dmFsdWUy"
         finally:
-            astra_db.delete_collection(collection_name)
+            retry(astra_db.delete_collection, args=[collection_name])
 
     def test_indexing_detection(self, astra_db: AstraDB) -> None:
         """Test the behaviour against preexisting legacy collections."""
@@ -197,6 +198,6 @@ def test_indexing_detection(self, astra_db: AstraDB) -> None:
                 collection_name="lc_test_custom_store", astra_db_client=astra_db
             )
 
-        astra_db.delete_collection("lc_test_legacy_store")
-        astra_db.delete_collection("lc_test_custom_store")
-        astra_db.delete_collection("lc_test_regular_store")
+        retry(astra_db.delete_collection, args=["lc_test_legacy_store"])
+        retry(astra_db.delete_collection, args=["lc_test_custom_store"])
+        retry(astra_db.delete_collection, args=["lc_test_regular_store"])