fix postgres test

HaoXuAI · HaoXuAI · commit d4f263929aed · 2024-04-14T15:34:53.000-07:00
diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py
@@ -9,6 +9,7 @@
 from psycopg2 import sql
 from psycopg2.extras import execute_values
 from psycopg2.pool import SimpleConnectionPool
+
 from feast import Entity
 from feast.feature_view import FeatureView
 from feast.infra.key_encoding_utils import serialize_entity_key
@@ -320,8 +321,8 @@ def retrieve_online_documents(
             for entity_key, feature_name, value, distance, event_ts in rows:
 
                 # TODO Deserialize entity_key to return the entity in response
-                entity_key_proto = EntityKeyProto()
-                entity_key_proto_bin = bytes(entity_key)
+                # entity_key_proto = EntityKeyProto()
+                # entity_key_proto_bin = bytes(entity_key)
 
                 # TODO Convert to List[float] for value type proto
                 feature_value_proto = ValueProto(string_val=value)
diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres_repo_configuration.py b/sdk/python/feast/infra/online_stores/contrib/postgres_repo_configuration.py
@@ -2,8 +2,8 @@
     IntegrationTestRepoConfig,
 )
 from tests.integration.feature_repos.universal.online_store.postgres import (
+    PGVectorOnlineStoreCreator,
     PostgresOnlineStoreCreator,
-    PGVectorOnlineStoreCreator
 )
 
 FULL_REPO_CONFIGS = [
diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py
@@ -26,7 +26,10 @@
 from feast.data_source import DataSource
 from feast.feature_store import FeatureStore  # noqa: E402
 from feast.wait import wait_retry_backoff  # noqa: E402
-from tests.data.data_creator import create_basic_driver_dataset, create_document_dataset  # noqa: E402
+from tests.data.data_creator import (  # noqa: E402
+    create_basic_driver_dataset,
+    create_document_dataset,
+)
 from tests.integration.feature_repos.integration_test_repo_config import (  # noqa: E402
     IntegrationTestRepoConfig,
 )
diff --git a/sdk/python/tests/data/data_creator.py b/sdk/python/tests/data/data_creator.py
@@ -95,4 +95,4 @@ def create_document_dataset() -> pd.DataFrame:
             pd.Timestamp(datetime.utcnow()).round("ms"),
         ],
     }
-    return pd.DataFrame(data)
+    return pd.DataFrame(data)
diff --git a/sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py b/sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py
@@ -54,7 +54,7 @@ def create_online_store(self) -> Dict[str, str]:
             timeout=10
         )
         command = "psql -h localhost -p 5432 -U root -d test -c 'CREATE EXTENSION IF NOT EXISTS vector;'"
-        res = self.container.exec(command)
+        self.container.exec(command)
         return {
             "host": "localhost",
             "type": "postgres",
@@ -67,4 +67,4 @@ def create_online_store(self) -> Dict[str, str]:
         }
 
     def teardown(self):
-        self.container.stop()
+        self.container.stop()
diff --git a/sdk/python/tests/integration/online_store/test_universal_online.py b/sdk/python/tests/integration/online_store/test_universal_online.py
@@ -28,8 +28,8 @@
 from tests.integration.feature_repos.universal.entities import driver, item
 from tests.integration.feature_repos.universal.feature_views import (
     create_driver_hourly_stats_feature_view,
+    create_item_embeddings_feature_view,
     driver_feature_view,
-    create_item_embeddings_feature_view
 )
 from tests.utils.data_source_test_creator import prep_file_source
 

Original file line number	Diff line number	Diff line change
`@@ -2,8 +2,8 @@`
`2`	`2`	`IntegrationTestRepoConfig,`
`3`	`3`	`)`
`4`	`4`	`from tests.integration.feature_repos.universal.online_store.postgres import (`
	`5`	`+ PGVectorOnlineStoreCreator,`
`5`	`6`	`PostgresOnlineStoreCreator,`
`6`		`- PGVectorOnlineStoreCreator`
`7`	`7`	`)`
`8`	`8`
`9`	`9`	`FULL_REPO_CONFIGS = [`
Original file line number	Diff line number	Diff line change
`@@ -95,4 +95,4 @@ def create_document_dataset() -> pd.DataFrame:`
`95`	`95`	`pd.Timestamp(datetime.utcnow()).round("ms"),`
`96`	`96`	`],`
`97`	`97`	`}`
`98`		`- return pd.DataFrame(data)`
	`98`	`+ return pd.DataFrame(data)`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ def create_online_store(self) -> Dict[str, str]:`
`54`	`54`	`timeout=10`
`55`	`55`	`)`
`56`	`56`	`command = "psql -h localhost -p 5432 -U root -d test -c 'CREATE EXTENSION IF NOT EXISTS vector;'"`
`57`		`- res = self.container.exec(command)`
	`57`	`+ self.container.exec(command)`
`58`	`58`	`return {`
`59`	`59`	`"host": "localhost",`
`60`	`60`	`"type": "postgres",`
`@@ -67,4 +67,4 @@ def create_online_store(self) -> Dict[str, str]:`
`67`	`67`	`}`
`68`	`68`
`69`	`69`	`def teardown(self):`
`70`		`- self.container.stop()`
	`70`	`+ self.container.stop()`
Original file line number	Diff line number	Diff line change
`@@ -28,8 +28,8 @@`
`28`	`28`	`from tests.integration.feature_repos.universal.entities import driver, item`
`29`	`29`	`from tests.integration.feature_repos.universal.feature_views import (`
`30`	`30`	`create_driver_hourly_stats_feature_view,`
	`31`	`+ create_item_embeddings_feature_view,`
`31`	`32`	`driver_feature_view,`
`32`		`- create_item_embeddings_feature_view`
`33`	`33`	`)`
`34`	`34`	`from tests.utils.data_source_test_creator import prep_file_source`
`35`	`35`