feat: Add tests to reach 90% coverage (#166)

* feat: Add tests to reach 90% coverage

* replace alloydb occurances

* add more tests

* fix key value error

* add is valid index test

* test

* set test orders

* delete skip python version

Author: duwenxin99

.coveragerc

@@ -5,4 +5,4 @@ omit =
 
 [report]
 show_missing = true
-fail_under = 82
+fail_under = 90

DEVELOPER.md

@@ -28,7 +28,7 @@ Learn more by reading [How should I write my commits?](https://github.com/google
 
 ### CI Platform Setup
 
-Cloud Build is used to run tests against Google Cloud resources in test project: langchain-alloydb-testing.
+Cloud Build is used to run tests against Google Cloud resources in test project: langchain-cloud-sql-testing.
 Each test has a corresponding Cloud Build trigger, see [all triggers][triggers].
 These tests are registered as required tests in `.github/sync-repo-settings.yaml`.
 
@@ -41,7 +41,7 @@ name: pg-integration-test-pr-py38
 description: Run integration tests on PR for Python 3.8
 filename: integration.cloudbuild.yaml
 github:
-  name: langchain-google-alloydb-pg-python
+  name: langchain-google-cloud-sql-pg-python
   owner: googleapis
   pullRequest:
     branch: .*

src/langchain_google_cloud_sql_pg/loader.py

@@ -555,29 +555,3 @@ def delete(self, docs: List[Document]) -> None:
             docs (List[langchain_core.documents.Document]): a list of documents to be deleted.
         """
         self.engine._run_as_sync(self.adelete(docs))
-
-    async def _aload_table_schema(self) -> sqlalchemy.Table:
-        """
-        Load table schema from existing table in PgSQL database.
-
-        Returns:
-            (sqlalchemy.Table): The loaded table.
-        """
-        metadata = sqlalchemy.MetaData()
-        async with self.engine._engine.connect() as conn:
-            await conn.run_sync(metadata.reflect, only=[self.table_name])
-
-        table = sqlalchemy.Table(self.table_name, metadata)
-        # Extract the schema information
-        schema = []
-        for column in table.columns:
-            schema.append(
-                {
-                    "name": column.name,
-                    "type": column.type.python_type,
-                    "max_length": getattr(column.type, "length", None),
-                    "nullable": not column.nullable,
-                }
-            )
-
-        return metadata.tables[self.table_name]

src/langchain_google_cloud_sql_pg/vectorstore.py

@@ -183,7 +183,7 @@ async def create(
             del all_columns[id_column]
             del all_columns[content_column]
             del all_columns[embedding_column]
-            metadata_columns = [k for k, _ in all_columns.keys()]
+            metadata_columns = [k for k in all_columns.keys()]
 
         return cls(
             cls.__create_key,

tests/test_cloudsql_vectorstore.py

@@ -133,6 +133,19 @@ async def vs_custom(self, engine):
         yield vs
         await engine._aexecute(f'DROP TABLE IF EXISTS "{CUSTOM_TABLE}"')
 
+    async def test_init_with_constructor(self, engine):
+        with pytest.raises(Exception):
+            PostgresVectorStore(
+                engine,
+                embedding_service=embeddings_service,
+                table_name=CUSTOM_TABLE,
+                id_column="myid",
+                content_column="noname",
+                embedding_column="myembedding",
+                metadata_columns=["page", "source"],
+                metadata_json_column="mymeta",
+            )
+
     async def test_post_init(self, engine):
         with pytest.raises(ValueError):
             await PostgresVectorStore.create(
@@ -265,4 +278,60 @@ async def test_add_texts(self, engine_sync, vs_sync):
         results = engine_sync._fetch(f'SELECT * FROM "{DEFAULT_TABLE_SYNC}"')
         assert len(results) == 6
 
+    async def test_ignore_metadata_columns(self, vs_custom):
+        column_to_ignore = "source"
+        vs = await PostgresVectorStore.create(
+            vs_custom.engine,
+            embedding_service=embeddings_service,
+            table_name=CUSTOM_TABLE,
+            ignore_metadata_columns=[column_to_ignore],
+            id_column="myid",
+            content_column="mycontent",
+            embedding_column="myembedding",
+            metadata_json_column="mymeta",
+        )
+        assert column_to_ignore not in vs.metadata_columns
+
+    async def test_create_vectorstore_with_invalid_parameters(self, vs_custom):
+        with pytest.raises(ValueError):
+            await PostgresVectorStore.create(
+                vs_custom.engine,
+                embedding_service=embeddings_service,
+                table_name=CUSTOM_TABLE,
+                id_column="myid",
+                content_column="mycontent",
+                embedding_column="myembedding",
+                metadata_columns=["random_column"],  # invalid metadata column
+            )
+        with pytest.raises(ValueError):
+            await PostgresVectorStore.create(
+                vs_custom.engine,
+                embedding_service=embeddings_service,
+                table_name=CUSTOM_TABLE,
+                id_column="myid",
+                content_column="langchain_id",  # invalid content column type
+                embedding_column="myembedding",
+                metadata_columns=["random_column"],
+            )
+        with pytest.raises(ValueError):
+            await PostgresVectorStore.create(
+                vs_custom.engine,
+                embedding_service=embeddings_service,
+                table_name=CUSTOM_TABLE,
+                id_column="myid",
+                content_column="mycontent",
+                embedding_column="random_column",  # invalid embedding column
+                metadata_columns=["random_column"],
+            )
+        with pytest.raises(ValueError):
+            await PostgresVectorStore.create(
+                vs_custom.engine,
+                embedding_service=embeddings_service,
+                table_name=CUSTOM_TABLE,
+                id_column="myid",
+                content_column="mycontent",
+                embedding_column="langchain_id",  # invalid embedding column data type
+                metadata_columns=["random_column"],
+            )
+
     # Need tests for store metadata=False

tests/test_cloudsql_vectorstore_index.py

@@ -55,10 +55,6 @@ def get_env_var(key: str, desc: str) -> str:
 
 
 @pytest.mark.asyncio(scope="class")
-@pytest.mark.skipif(
-    sys.version_info != (3, 11),
-    reason="To prevent index clashes only run on python3.11 or higher",
-)
 class TestIndex:
     @pytest.fixture(scope="module")
     def db_project(self) -> str:
@@ -101,11 +97,13 @@ async def vs(self, engine):
         await engine._aexecute(f"DROP TABLE IF EXISTS {DEFAULT_TABLE}")
         await engine._engine.dispose()
 
+    @pytest.mark.run(order=1)
     async def test_aapply_vector_index(self, vs):
         index = HNSWIndex()
         await vs.aapply_vector_index(index)
         assert await vs.is_valid_index(DEFAULT_INDEX_NAME)
 
+    @pytest.mark.run(order=2)
     async def test_areindex(self, vs):
         if not await vs.is_valid_index(DEFAULT_INDEX_NAME):
             index = HNSWIndex()
@@ -114,6 +112,7 @@ async def test_areindex(self, vs):
         await vs.areindex(DEFAULT_INDEX_NAME)
         assert await vs.is_valid_index(DEFAULT_INDEX_NAME)
 
+    @pytest.mark.run(order=3)
     async def test_dropindex(self, vs):
         await vs.adrop_vector_index()
         result = await vs.is_valid_index(DEFAULT_INDEX_NAME)
@@ -130,3 +129,7 @@ async def test_aapply_vector_index_ivfflat(self, vs):
         await vs.aapply_vector_index(index)
         assert await vs.is_valid_index("secondindex")
         await vs.adrop_vector_index("secondindex")
+
+    async def test_is_valid_index(self, vs):
+        is_valid = await vs.is_valid_index("invalid_index")
+        assert is_valid == False

tests/test_cloudsql_vectorstore_search.py

@@ -21,7 +21,7 @@
 from langchain_core.documents import Document
 
 from langchain_google_cloud_sql_pg import Column, PostgresEngine, PostgresVectorStore
-from langchain_google_cloud_sql_pg.indexes import HNSWQueryOptions, IVFFlatQueryOptions
+from langchain_google_cloud_sql_pg.indexes import DistanceStrategy, HNSWQueryOptions
 
 DEFAULT_TABLE = "test_table" + str(uuid.uuid4()).replace("-", "_")
 CUSTOM_TABLE = "test_table_custom" + str(uuid.uuid4()).replace("-", "_")
@@ -151,7 +151,7 @@ async def test_asimilarity_search_by_vector(self, vs):
         assert results[0][0] == Document(page_content="foo")
         assert results[0][1] == 0
 
-    async def test_similarity_search_with_relevance_scores_threshold(self, vs):
+    async def test_similarity_search_with_relevance_scores_threshold_cosine(self, vs):
         score_threshold = {"score_threshold": 0}
         results = await vs.asimilarity_search_with_relevance_scores(
             "foo", **score_threshold
@@ -171,6 +171,23 @@ async def test_similarity_search_with_relevance_scores_threshold(self, vs):
         assert len(results) == 1
         assert results[0][0] == Document(page_content="foo")
 
+    async def test_similarity_search_with_relevance_scores_threshold_euclidean(
+        self, engine
+    ):
+        vs = await PostgresVectorStore.create(
+            engine,
+            embedding_service=embeddings_service,
+            table_name=DEFAULT_TABLE,
+            distance_strategy=DistanceStrategy.EUCLIDEAN,
+        )
+
+        score_threshold = {"score_threshold": 0.9}
+        results = await vs.asimilarity_search_with_relevance_scores(
+            "foo", **score_threshold
+        )
+        assert len(results) == 1
+        assert results[0][0] == Document(page_content="foo")
+
     async def test_amax_marginal_relevance_search(self, vs):
         results = await vs.amax_marginal_relevance_search("bar")
         assert results[0] == Document(page_content="bar")

tests/test_postgresql_loader.py

@@ -69,6 +69,29 @@ async def _cleanup_table(self, engine):
         query = f'DROP TABLE IF EXISTS "{table_name}"'
         await engine._aexecute(query)
 
+    async def test_create_loader_with_invalid_parameters(self, engine):
+        with pytest.raises(ValueError):
+            await PostgresLoader.create(
+                engine=engine,
+            )
+        with pytest.raises(ValueError):
+
+            def fake_formatter():
+                return None
+
+            await PostgresLoader.create(
+                engine=engine,
+                table_name=table_name,
+                format="text",
+                formatter=fake_formatter,
+            )
+        with pytest.raises(ValueError):
+            await PostgresLoader.create(
+                engine=engine,
+                table_name=table_name,
+                format="fake_format",
+            )
+
     async def test_load_from_query_default(self, engine):
         try:
             await self._cleanup_table(engine)
@@ -216,6 +239,30 @@ async def test_load_from_query_customized_content_default_metadata(self, engine)
                 )
             ]
 
+            loader = await PostgresLoader.create(
+                engine=engine,
+                query=f'SELECT * FROM "{table_name}";',
+                content_columns=[
+                    "variety",
+                    "quantity_in_stock",
+                    "price_per_unit",
+                ],
+                format="JSON",
+            )
+
+            documents = await self._collect_async_items(loader.alazy_load())
+
+            assert documents == [
+                Document(
+                    page_content='{"variety": "Granny Smith", "quantity_in_stock": 150, "price_per_unit": 1}',
+                    metadata={
+                        "fruit_id": 1,
+                        "fruit_name": "Apple",
+                        "organic": 1,
+                    },
+                )
+            ]
+
         finally:
             await self._cleanup_table(engine)