From 752f957669a139fa92cac3d2258dcd6e4500d3fc Mon Sep 17 00:00:00 2001
From: Jorge <jzaldivar@google.com>
Date: Fri, 22 Mar 2024 00:35:57 +0100
Subject: [PATCH] Add integration tests

---
 .../integration_tests/test_vectorstores.py    | 47 ++++++++++++-------
 1 file changed, 30 insertions(+), 17 deletions(-)

diff --git a/libs/vertexai/tests/integration_tests/test_vectorstores.py b/libs/vertexai/tests/integration_tests/test_vectorstores.py
index c8040a3d..72cf88c0 100644
--- a/libs/vertexai/tests/integration_tests/test_vectorstores.py
+++ b/libs/vertexai/tests/integration_tests/test_vectorstores.py
@@ -66,6 +66,7 @@ def datastore_document_storage(
 
 @pytest.fixture
 def vector_store() -> VectorSearchVectorStore:
+
     embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default")
 
     vector_store = VectorSearchVectorStore.from_components(
@@ -82,15 +83,16 @@ def vector_store() -> VectorSearchVectorStore:
 
 @pytest.fixture
 def datastore_vector_store() -> VectorSearchVectorStoreDatastore:
+
     embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default")
 
     vector_store = VectorSearchVectorStoreDatastore.from_components(
         project_id=os.environ["PROJECT_ID"],
         region=os.environ["REGION"],
-        index_staging_bucket_name=os.environ["GCS_BUCKET_NAME"],
-        index_id=os.environ["INDEX_ID"],
-        endpoint_id=os.environ["ENDPOINT_ID"],
+        index_id=os.environ["STREAM_INDEX_ID_DATASTORE"],
+        endpoint_id=os.environ["STREAM_ENDPOINT_ID_DATASTORE"],
         embedding=embeddings,
+        stream_update=True
     )
 
     return vector_store
@@ -164,8 +166,12 @@ def test_public_endpoint_vector_searcher(sdk_manager: VectorSearchSDKManager):
 
 
 @pytest.mark.extended
-def test_vector_store(vector_store: VectorSearchVectorStore):
-    assert isinstance(vector_store, VectorSearchVectorStore)
+@pytest.mark.parametrize(
+    "vector_store_class", ["vector_store", "datastore_vector_store"]
+)
+def test_vector_store(vector_store_class: str, request: pytest.FixtureRequest):
+
+    vector_store: VectorSearchVectorStore = request.getfixturevalue(vector_store_class)
 
     query = "What are your favourite animals?"
     docs_with_scores = vector_store.similarity_search_with_score(query, k=1)
@@ -181,7 +187,16 @@ def test_vector_store(vector_store: VectorSearchVectorStore):
 
 
 @pytest.mark.extended
-def test_vector_store_filtering(vector_store: VectorSearchVectorStore):
+@pytest.mark.parametrize(
+    "vector_store_class", [
+        "vector_store", 
+        #"datastore_vector_store" Waiting for the bug to be fixed as its stream
+    ]
+)
+def test_vector_store_filtering(
+    vector_store_class: str, request: pytest.FixtureRequest):
+
+    vector_store: VectorSearchVectorStore = request.getfixturevalue(vector_store_class)
     documents = vector_store.similarity_search(
         "I want some pants",
         filter=[Namespace(name="color", allow_tokens=["blue"])],
@@ -194,19 +209,17 @@ def test_vector_store_filtering(vector_store: VectorSearchVectorStore):
 
 
 @pytest.mark.extended
-def test_vector_store_update_index(sample_documents: List[Document]):
-    embeddings = VertexAIEmbeddings(model_name="textembedding-gecko-default")
+def test_vector_store_update_index(
+    vector_store: VectorSearchVectorStore, sample_documents: List[Document]):
+    vector_store.add_documents(documents=sample_documents, is_complete_overwrite=True)
 
-    vector_store = VectorSearchVectorStore.from_components(
-        project_id=os.environ["PROJECT_ID"],
-        region=os.environ["REGION"],
-        gcs_bucket_name=os.environ["GCS_BUCKET_NAME"],
-        index_id=os.environ["INDEX_ID"],
-        endpoint_id=os.environ["ENDPOINT_ID"],
-        embedding=embeddings,
-    )
+@pytest.mark.extended
+def test_vector_store_stream_update_index(
+    datastore_vector_store: VectorSearchVectorStoreDatastore, 
+    sample_documents: List[Document]):
 
-    vector_store.add_documents(documents=sample_documents, is_complete_overwrite=True)
+    datastore_vector_store.add_documents(
+        documents=sample_documents, is_complete_overwrite=True)
 
 
 @pytest.fixture