chore: Add fixture so that benchmark test is also tested in integration tests (#3065)

Add fixture so that benchmark test is also tested in integration tests

Signed-off-by: Felix Wang <wangfelix98@gmail.com>

Signed-off-by: Felix Wang <wangfelix98@gmail.com>

Author: felixwang9817

sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py

@@ -1,64 +1,14 @@
-import random
-from typing import List
-
 import pytest
 
-from feast import FeatureService
-from feast.feast_object import FeastObject
-from tests.integration.feature_repos.repo_configuration import (
-    construct_universal_feature_views,
-)
-from tests.integration.feature_repos.universal.entities import (
-    customer,
-    driver,
-    location,
-)
-
 
 @pytest.mark.benchmark
 @pytest.mark.integration
 @pytest.mark.universal_online_stores
-def test_online_retrieval(environment, universal_data_sources, benchmark):
-    fs = environment.feature_store
-    entities, datasets, data_sources = universal_data_sources
-    feature_views = construct_universal_feature_views(data_sources)
-
-    feature_service = FeatureService(
-        name="convrate_plus100",
-        features=[feature_views.driver[["conv_rate"]], feature_views.driver_odfv],
-    )
-
-    feast_objects: List[FeastObject] = []
-    feast_objects.extend(feature_views.values())
-    feast_objects.extend([driver(), customer(), location(), feature_service])
-    fs.apply(feast_objects)
-    fs.materialize(environment.start_date, environment.end_date)
-
-    sample_drivers = random.sample(entities.driver_vals, 10)
-
-    sample_customers = random.sample(entities.customer_vals, 10)
-
-    entity_rows = [
-        {"driver_id": d, "customer_id": c, "val_to_add": 50}
-        for (d, c) in zip(sample_drivers, sample_customers)
-    ]
-
-    feature_refs = [
-        "driver_stats:conv_rate",
-        "driver_stats:avg_daily_trips",
-        "customer_profile:current_balance",
-        "customer_profile:avg_passenger_count",
-        "customer_profile:lifetime_trip_count",
-        "conv_rate_plus_100:conv_rate_plus_100",
-        "conv_rate_plus_100:conv_rate_plus_val_to_add",
-        "global_stats:num_rides",
-        "global_stats:avg_ride_length",
-    ]
-    unprefixed_feature_refs = [f.rsplit(":", 1)[-1] for f in feature_refs if ":" in f]
-    # Remove the on demand feature view output features, since they're not present in the source dataframe
-    unprefixed_feature_refs.remove("conv_rate_plus_100")
-    unprefixed_feature_refs.remove("conv_rate_plus_val_to_add")
-
+def test_online_retrieval(feature_store_for_online_retrieval, benchmark):
+    """
+    Benchmarks a basic online retrieval flow.
+    """
+    fs, feature_refs, entity_rows = feature_store_for_online_retrieval
     benchmark(
         fs.get_online_features,
         features=feature_refs,

sdk/python/tests/conftest.py

@@ -14,18 +14,19 @@
 import logging
 import multiprocessing
 import os
+import random
 from datetime import datetime, timedelta
 from multiprocessing import Process
 from sys import platform
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Tuple
 
 import pandas as pd
 import pytest
 from _pytest.nodes import Item
 
 os.environ["FEAST_USAGE"] = "False"
 os.environ["IS_TEST"] = "True"
-from feast import FeatureStore  # noqa: E402
+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  # noqa: E402
 from tests.integration.feature_repos.integration_test_repo_config import (  # noqa: E402
@@ -38,11 +39,17 @@
     Environment,
     TestData,
     construct_test_environment,
+    construct_universal_feature_views,
     construct_universal_test_data,
 )
 from tests.integration.feature_repos.universal.data_sources.file import (  # noqa: E402
     FileDataSourceCreator,
 )
+from tests.integration.feature_repos.universal.entities import (  # noqa: E402
+    customer,
+    driver,
+    location,
+)
 from tests.utils.http_server import check_port_open, free_port  # noqa: E402
 
 logger = logging.getLogger(__name__)
@@ -373,3 +380,44 @@ def e2e_data_sources(environment: Environment):
     )
 
     return df, data_source
+
+
+@pytest.fixture
+def feature_store_for_online_retrieval(
+    environment, universal_data_sources
+) -> Tuple[FeatureStore, List[str], List[Dict[str, int]]]:
+    """
+    Returns a feature store that is ready for online retrieval, along with entity rows and feature
+    refs that can be used to query for online features.
+    """
+    fs = environment.feature_store
+    entities, datasets, data_sources = universal_data_sources
+    feature_views = construct_universal_feature_views(data_sources)
+
+    feast_objects = []
+    feast_objects.extend(feature_views.values())
+    feast_objects.extend([driver(), customer(), location()])
+    fs.apply(feast_objects)
+    fs.materialize(environment.start_date, environment.end_date)
+
+    sample_drivers = random.sample(entities.driver_vals, 10)
+    sample_customers = random.sample(entities.customer_vals, 10)
+
+    entity_rows = [
+        {"driver_id": d, "customer_id": c, "val_to_add": 50}
+        for (d, c) in zip(sample_drivers, sample_customers)
+    ]
+
+    feature_refs = [
+        "driver_stats:conv_rate",
+        "driver_stats:avg_daily_trips",
+        "customer_profile:current_balance",
+        "customer_profile:avg_passenger_count",
+        "customer_profile:lifetime_trip_count",
+        "conv_rate_plus_100:conv_rate_plus_100",
+        "conv_rate_plus_100:conv_rate_plus_val_to_add",
+        "global_stats:num_rides",
+        "global_stats:avg_ride_length",
+    ]
+
+    return fs, feature_refs, entity_rows

sdk/python/tests/integration/online_store/test_universal_online.py

@@ -13,11 +13,14 @@
 import requests
 from botocore.exceptions import BotoCoreError
 
-from feast import Entity, FeatureService, FeatureView, Field
+from feast.entity import Entity
 from feast.errors import (
     FeatureNameCollisionError,
     RequestDataNotFoundInEntityRowsException,
 )
+from feast.feature_service import FeatureService
+from feast.feature_view import FeatureView
+from feast.field import Field
 from feast.online_response import TIMESTAMP_POSTFIX
 from feast.types import Float32, Int32, String
 from feast.wait import wait_retry_backoff
@@ -767,6 +770,21 @@ def eventually_apply() -> Tuple[None, bool]:
     assert all(v is None for v in online_features["value"])
 
 
+@pytest.mark.integration
+@pytest.mark.universal_online_stores
+def test_online_retrieval_success(feature_store_for_online_retrieval):
+    """
+    Tests that online retrieval executes successfully (i.e. without errors).
+
+    Does not test for correctness of the results of online retrieval.
+    """
+    fs, feature_refs, entity_rows = feature_store_for_online_retrieval
+    fs.get_online_features(
+        features=feature_refs,
+        entity_rows=entity_rows,
+    )
+
+
 def response_feature_name(
     feature: str, feature_refs: List[str], full_feature_names: bool
 ) -> str: