Ensure AstraDB integration tests clean up the environment (#13774)

- **Description:** currently astra_db integration tests might leave
orphan collections
  - **Tag maintainer:** @hemidactylus 
  - **Twitter handle:** nicoloboschi
This commit is contained in:
Nicolò Boschi 2023-12-03 20:14:42 +01:00 committed by GitHub
parent 7bc4c12477
commit 0de7cf898d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -160,7 +160,9 @@ class TestAstraDB:
api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
)
try:
assert v_store.similarity_search("Ho", k=1)[0].page_content == "Ho"
finally:
v_store.delete_collection()
# from_texts
@ -175,8 +177,9 @@ class TestAstraDB:
api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
)
try:
assert v_store_2.similarity_search("Hoi", k=1)[0].page_content == "Hoi"
# manual collection delete
finally:
v_store_2.delete_collection()
def test_astradb_vectorstore_crud(self, store_someemb: AstraDB) -> None:
@ -355,10 +358,11 @@ class TestAstraDB:
def test_astradb_vectorstore_drop(self) -> None:
"""behaviour of 'delete_collection'."""
collection_name = "lc_test_d"
emb = SomeEmbeddings(dimension=2)
v_store = AstraDB(
embedding=emb,
collection_name="lc_test_d",
collection_name=collection_name,
token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
@ -368,7 +372,7 @@ class TestAstraDB:
# another instance pointing to the same collection on DB
v_store_kenny = AstraDB(
embedding=emb,
collection_name="lc_test_d",
collection_name=collection_name,
token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
@ -392,6 +396,7 @@ class TestAstraDB:
bulk_insert_overwrite_concurrency=7,
bulk_delete_concurrency=19,
)
try:
# add_texts
N = 50
texts = [str(i + 1 / 7.0) for i in range(N)]
@ -408,6 +413,7 @@ class TestAstraDB:
_ = v_store.delete(ids[: N // 2])
_ = v_store.delete(ids[N // 2 :], concurrency=23)
#
finally:
v_store.delete_collection()
def test_astradb_vectorstore_metrics(self) -> None:
@ -440,6 +446,7 @@ class TestAstraDB:
namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
metric="cosine",
)
try:
vstore_cos.add_texts(
texts=texts,
ids=ids,
@ -449,8 +456,10 @@ class TestAstraDB:
k=1,
)[0]
assert id_from_cos == "scaled"
finally:
vstore_cos.delete_collection()
# creation, population, query - euclidean
vstore_euc = AstraDB(
embedding=emb,
collection_name="lc_test_m_e",
@ -459,6 +468,7 @@ class TestAstraDB:
namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
metric="euclidean",
)
try:
vstore_euc.add_texts(
texts=texts,
ids=ids,
@ -468,4 +478,5 @@ class TestAstraDB:
k=1,
)[0]
assert id_from_euc == "rotated"
finally:
vstore_euc.delete_collection()