chore: Make contrib test pluggable (#2654)

* make contrib test plugable

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* no autouse

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* pass request

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* fix fixture name

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* publish trino fixture on package level

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* format

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* disable some tests for contrib tests run

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

* address comments

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

Makefile

@@ -72,11 +72,25 @@ test-python-integration-container:
 	FEAST_USAGE=False IS_TEST=True FEAST_LOCAL_ONLINE_CONTAINER=True python -m pytest -n 8 --integration sdk/python/tests
 
 test-python-universal-contrib:
-	PYTHONPATH='.' FULL_REPO_CONFIGS_MODULE=sdk.python.feast.infra.offline_stores.contrib.contrib_repo_configuration FEAST_USAGE=False IS_TEST=True python -m pytest -n 8 --integration --universal sdk/python/tests
+	PYTHONPATH='.' \
+	FULL_REPO_CONFIGS_MODULE=sdk.python.feast.infra.offline_stores.contrib.contrib_repo_configuration \
+	PYTEST_PLUGINS=feast.infra.offline_stores.contrib.trino_offline_store.tests \
+ 	FEAST_USAGE=False IS_TEST=True \
+ 	python -m pytest -n 8 --integration --universal \
+ 	 	-k "not test_historical_retrieval_fails_on_validation and \
+			not test_historical_retrieval_with_validation and \
+			not test_historical_features_persisting and \
+			not test_historical_retrieval_fails_on_validation and \
+			not test_universal_cli and \
+			not test_go_feature_server and \
+			not test_feature_logging and \
+			not test_universal_types" \
+ 	 sdk/python/tests
 
 test-python-universal-postgres:
 	PYTHONPATH='.' \
 		FULL_REPO_CONFIGS_MODULE=sdk.python.feast.infra.offline_stores.contrib.postgres_repo_configuration \
+		PYTEST_PLUGINS=sdk.python.feast.infra.offline_stores.contrib.postgres_offline_store.tests \
 		FEAST_USAGE=False \
 		IS_TEST=True \
 		python -m pytest -x --integration --universal \

sdk/python/feast/infra/offline_stores/contrib/contrib_repo_configuration.py

@@ -1,12 +1,12 @@
-from tests.integration.feature_repos.integration_test_repo_config import (
-    IntegrationTestRepoConfig,
-)
-from tests.integration.feature_repos.universal.data_sources.spark_data_source_creator import (
+from feast.infra.offline_stores.contrib.spark_offline_store.tests.data_source import (
     SparkDataSourceCreator,
 )
-from tests.integration.feature_repos.universal.data_sources.trino import (
+from feast.infra.offline_stores.contrib.trino_offline_store.tests.data_source import (
     TrinoSourceCreator,
 )
+from tests.integration.feature_repos.integration_test_repo_config import (
+    IntegrationTestRepoConfig,
+)
 
 FULL_REPO_CONFIGS = [
     IntegrationTestRepoConfig(offline_store_creator=SparkDataSourceCreator),

sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/tests/__init__.py

@@ -0,0 +1 @@
+from .data_source import postgres_container  # noqa

sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/tests/data_source.py

@@ -0,0 +1,122 @@
+import logging
+from typing import Dict, Optional
+
+import pandas as pd
+import pytest
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+
+from feast.data_source import DataSource
+from feast.infra.offline_stores.contrib.postgres_offline_store.postgres import (
+    PostgreSQLOfflineStoreConfig,
+    PostgreSQLSource,
+)
+from feast.infra.utils.postgres.connection_utils import df_to_postgres_table
+from tests.integration.feature_repos.universal.data_source_creator import (
+    DataSourceCreator,
+)
+from tests.integration.feature_repos.universal.online_store_creator import (
+    OnlineStoreCreator,
+)
+
+logger = logging.getLogger(__name__)
+
+POSTGRES_USER = "test"
+POSTGRES_PASSWORD = "test"
+POSTGRES_DB = "test"
+
+
+@pytest.fixture(scope="session")
+def postgres_container():
+    container = (
+        DockerContainer("postgres:latest")
+        .with_exposed_ports(5432)
+        .with_env("POSTGRES_USER", POSTGRES_USER)
+        .with_env("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
+        .with_env("POSTGRES_DB", POSTGRES_DB)
+    )
+
+    container.start()
+
+    log_string_to_wait_for = "database system is ready to accept connections"
+    waited = wait_for_logs(
+        container=container, predicate=log_string_to_wait_for, timeout=30, interval=10,
+    )
+    logger.info("Waited for %s seconds until postgres container was up", waited)
+
+    yield container
+    container.stop()
+
+
+class PostgreSQLDataSourceCreator(DataSourceCreator, OnlineStoreCreator):
+    def __init__(
+        self, project_name: str, fixture_request: pytest.FixtureRequest, **kwargs
+    ):
+        super().__init__(project_name,)
+
+        self.project_name = project_name
+        self.container = fixture_request.getfixturevalue("postgres_container")
+        if not self.container:
+            raise RuntimeError(
+                "In order to use this data source "
+                "'feast.infra.offline_stores.contrib.postgres_offline_store.tests' "
+                "must be include into pytest plugins"
+            )
+
+        self.offline_store_config = PostgreSQLOfflineStoreConfig(
+            type="postgres",
+            host="localhost",
+            port=self.container.get_exposed_port(5432),
+            database=self.container.env["POSTGRES_DB"],
+            db_schema="public",
+            user=self.container.env["POSTGRES_USER"],
+            password=self.container.env["POSTGRES_PASSWORD"],
+        )
+
+    def create_data_source(
+        self,
+        df: pd.DataFrame,
+        destination_name: str,
+        suffix: Optional[str] = None,
+        timestamp_field="ts",
+        created_timestamp_column="created_ts",
+        field_mapping: Dict[str, str] = None,
+    ) -> DataSource:
+        destination_name = self.get_prefixed_table_name(destination_name)
+
+        if self.offline_store_config:
+            df_to_postgres_table(self.offline_store_config, df, destination_name)
+
+        return PostgreSQLSource(
+            name=destination_name,
+            query=f"SELECT * FROM {destination_name}",
+            timestamp_field=timestamp_field,
+            created_timestamp_column=created_timestamp_column,
+            field_mapping=field_mapping or {"ts_1": "ts"},
+        )
+
+    def create_offline_store_config(self) -> PostgreSQLOfflineStoreConfig:
+        assert self.offline_store_config
+        return self.offline_store_config
+
+    def get_prefixed_table_name(self, suffix: str) -> str:
+        return f"{self.project_name}_{suffix}"
+
+    def create_online_store(self) -> Dict[str, str]:
+        assert self.container
+        return {
+            "type": "postgres",
+            "host": "localhost",
+            "port": self.container.get_exposed_port(5432),
+            "database": POSTGRES_DB,
+            "db_schema": "feature_store",
+            "user": POSTGRES_USER,
+            "password": POSTGRES_PASSWORD,
+        }
+
+    def create_saved_dataset_destination(self):
+        # FIXME: ...
+        return None
+
+    def teardown(self):
+        pass

sdk/python/feast/infra/offline_stores/contrib/postgres_repo_configuration.py

@@ -1,9 +1,9 @@
+from feast.infra.offline_stores.contrib.postgres_offline_store.tests.data_source import (
+    PostgreSQLDataSourceCreator,
+)
 from tests.integration.feature_repos.integration_test_repo_config import (
     IntegrationTestRepoConfig,
 )
-from tests.integration.feature_repos.universal.data_sources.postgres import (
-    PostgreSQLDataSourceCreator,
-)
 
 FULL_REPO_CONFIGS = [
     IntegrationTestRepoConfig(

sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/test_config/manual_tests.py

@@ -1,9 +1,9 @@
+from feast.infra.offline_stores.contrib.trino_offline_store.tests.data_source import (
+    TrinoSourceCreator,
+)
 from tests.integration.feature_repos.integration_test_repo_config import (
     IntegrationTestRepoConfig,
 )
-from tests.integration.feature_repos.universal.data_sources.trino import (
-    TrinoSourceCreator,
-)
 
 FULL_REPO_CONFIGS = [
     IntegrationTestRepoConfig(

sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/__init__.py

@@ -0,0 +1 @@
+from .data_source import trino_container  # noqa

sdk/python/tests/integration/feature_repos/universal/data_sources/trino.py → sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/data_source.py

@@ -3,6 +3,7 @@
 from typing import Dict, List, Optional
 
 import pandas as pd
+import pytest
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
@@ -24,46 +25,49 @@
 )
 
 
+@pytest.fixture(scope="session")
+def trino_container():
+    current_file = pathlib.Path(__file__).parent.resolve()
+    catalog_dir = current_file.parent.joinpath("catalog")
+    container = (
+        DockerContainer("trinodb/trino:376")
+        .with_volume_mapping(catalog_dir, "/etc/catalog/")
+        .with_exposed_ports("8080")
+    )
+
+    container.start()
+
+    log_string_to_wait_for = "SERVER STARTED"
+    wait_for_logs(container=container, predicate=log_string_to_wait_for, timeout=30)
+
+    yield container
+
+    container.stop()
+
+
 class TrinoSourceCreator(DataSourceCreator):
 
     tables: List[str] = []
 
-    def __init__(self, project_name: str, **kwargs):
+    def __init__(
+        self, project_name: str, fixture_request: pytest.FixtureRequest, **kwargs
+    ):
         super().__init__(project_name)
         self.tables_created: List[str] = []
-
-        if "offline_container" not in kwargs or not kwargs.get(
-            "offline_container", None
-        ):
-            # If we don't get an offline container provided, we try to create it on the fly.
-            # the problem here is that each test creates its own conatiner, which basically
-            # browns out developer laptops.
-            current_file = pathlib.Path(__file__).parent.resolve()
-            catalog_dir = current_file.parent.joinpath("catalog")
-            self.container = (
-                DockerContainer("trinodb/trino:376")
-                .with_volume_mapping(catalog_dir, "/etc/catalog/")
-                .with_exposed_ports("8080")
+        self.container = fixture_request.getfixturevalue("trino_container")
+        if not self.container:
+            raise RuntimeError(
+                "In order to use this data source "
+                "'feast.infra.offline_stores.contrib.trino_offline_store.tests' "
+                "must be include into pytest plugins"
             )
-
-            self.container.start()
-            self.provided_container = False
-            log_string_to_wait_for = "SERVER STARTED"
-            wait_for_logs(
-                container=self.container, predicate=log_string_to_wait_for, timeout=30
-            )
-        else:
-            self.provided_container = True
-            self.container = kwargs["offline_container"]
-
         self.exposed_port = self.container.get_exposed_port("8080")
         self.client = Trino(
             user="user", catalog="memory", host="localhost", port=self.exposed_port,
         )
 
     def teardown(self):
-        if not self.provided_container:
-            self.container.stop()
+        pass
 
     def create_data_source(
         self,

sdk/python/feast/infra/online_stores/contrib/postgres.py

@@ -131,7 +131,7 @@ def online_read(
                     res = {}
                     for feature_name, value_bin, event_ts in value:
                         val = ValueProto()
-                        val.ParseFromString(value_bin)
+                        val.ParseFromString(bytes(value_bin))
                         res[feature_name] = val
                     result.append((event_ts, res))
                 else: