fix the apikey header-precedence-over-kms tests

langchain-ai · Sep 30, 2024 · 52e5415 · 52e5415
1 parent f026755
commit 52e5415
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/libs/astradb/tests/integration_tests/test_vectorstore_ddl_tests.py b/libs/astradb/tests/integration_tests/test_vectorstore_ddl_tests.py
@@ -498,11 +498,12 @@ def test_astradb_vectorstore_vectorize_headers_precedence_stringheader(
             api_endpoint=astra_db_credentials["api_endpoint"],
             namespace=astra_db_credentials["namespace"],
             environment=astra_db_credentials["environment"],
-            setup_mode=SetupMode.OFF,
             collection_vector_service_options=OPENAI_VECTORIZE_OPTIONS_KMS,
             collection_embedding_api_key="verywrong",
         )
-        with pytest.raises(InsertManyException):
+        # More specific messages are provider-specific, such as OpenAI returning:
+        # "... Incorrect API key provided: verywrong ..."
+        with pytest.raises(InsertManyException, match="Embedding Provider returned"):
             v_store.add_texts(["Failing"])
 
     @pytest.mark.skipif(
@@ -525,9 +526,10 @@ def test_astradb_vectorstore_vectorize_headers_precedence_headerprovider(
             api_endpoint=astra_db_credentials["api_endpoint"],
             namespace=astra_db_credentials["namespace"],
             environment=astra_db_credentials["environment"],
-            setup_mode=SetupMode.OFF,
             collection_vector_service_options=OPENAI_VECTORIZE_OPTIONS_KMS,
             collection_embedding_api_key=EmbeddingAPIKeyHeaderProvider("verywrong"),
         )
-        with pytest.raises(InsertManyException):
+        # More specific messages are provider-specific, such as OpenAI returning:
+        # "... Incorrect API key provided: verywrong ..."
+        with pytest.raises(InsertManyException, match="Embedding Provider returned"):
             v_store.add_texts(["Failing"])