fixing tests

deepset-ai · Feb 13, 2024 · 7b87fe0 · 7b87fe0
1 parent cdeb8b1
commit 7b87fe0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py
@@ -118,9 +118,7 @@ def __init__(
             https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw
         """
 
-        self.connection_string = (
-            connection_string if isinstance(connection_string, str) else connection_string.resolve_value()
-        )
+        self.connection_string = connection_string
         self.table_name = table_name
         self.embedding_dimension = embedding_dimension
         if vector_function not in VALID_VECTOR_FUNCTIONS:
@@ -133,7 +131,7 @@ def __init__(
         self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {}
         self.hnsw_ef_search = hnsw_ef_search
 
-        connection = connect(self.connection_string)
+        connection = connect(self.connection_string.resolve_value())
         connection.autocommit = True
         self._connection = connection
 
@@ -154,7 +152,7 @@ def __init__(
     def to_dict(self) -> Dict[str, Any]:
         return default_to_dict(
             self,
-            connection_string=self.connection_string,
+            connection_string=self.connection_string.to_dict(),
             table_name=self.table_name,
             embedding_dimension=self.embedding_dimension,
             vector_function=self.vector_function,

diff --git a/integrations/pgvector/tests/test_document_store.py b/integrations/pgvector/tests/test_document_store.py
@@ -51,7 +51,6 @@ def test_init(self):
             hnsw_ef_search=50,
         )
 
-        assert document_store.connection_string == "postgresql://postgres:postgres@localhost:5432/postgres"
         assert document_store.table_name == "my_table"
         assert document_store.embedding_dimension == 512
         assert document_store.vector_function == "l2_distance"
@@ -76,7 +75,7 @@ def test_to_dict(self):
         assert document_store.to_dict() == {
             "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
             "init_parameters": {
-                "connection_string": "postgresql://postgres:postgres@localhost:5432/postgres",
+                "connection_string": {'env_vars': ['PG_CONN_STR'], 'strict': True, 'type': 'env_var'},
                 "table_name": "my_table",
                 "embedding_dimension": 512,
                 "vector_function": "l2_distance",