diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py index a095255d5af..8593512d5c6 100644 --- a/sdk/python/feast/feature_server.py +++ b/sdk/python/feast/feature_server.py @@ -6,7 +6,7 @@ import traceback from contextlib import asynccontextmanager from importlib import resources as importlib_resources -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union import pandas as pd import psutil @@ -86,10 +86,18 @@ class MaterializeIncrementalRequest(BaseModel): class GetOnlineFeaturesRequest(BaseModel): entities: Dict[str, List[Any]] feature_service: Optional[str] = None - features: Optional[List[str]] = None + features: List[str] = [] full_feature_names: bool = False - query_embedding: Optional[List[float]] = None + + +class GetOnlineDocumentsRequest(BaseModel): + feature_service: Optional[str] = None + features: List[str] = [] + full_feature_names: bool = False + top_k: Optional[int] = None + query: Optional[List[float]] = None query_string: Optional[str] = None + api_version: Optional[int] = 1 class ChatMessage(BaseModel): @@ -110,7 +118,10 @@ class SaveDocumentRequest(BaseModel): data: dict -def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"): +def _get_features( + request: Union[GetOnlineFeaturesRequest, GetOnlineDocumentsRequest], + store: "feast.FeatureStore", +): if request.feature_service: feature_service = store.get_feature_service( request.feature_service, allow_cache=True @@ -246,7 +257,7 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, An dependencies=[Depends(inject_user_details)], ) async def retrieve_online_documents( - request: GetOnlineFeaturesRequest, + request: GetOnlineDocumentsRequest, ) -> Dict[str, Any]: logger.warning( "This endpoint is in alpha and will be moved to /get-online-features when stable." @@ -254,16 +265,18 @@ async def retrieve_online_documents( # Initialize parameters for FeatureStore.retrieve_online_documents_v2(...) call features = await run_in_threadpool(_get_features, request, store) - read_params = dict( - features=features, - full_feature_names=request.full_feature_names, - query=request.query_embedding, - query_string=request.query_string, - ) + read_params = dict(features=features, query=request.query, top_k=request.top_k) + if request.api_version == 2 and request.query_string is not None: + read_params["query_string"] = request.query_string - response = await run_in_threadpool( - lambda: store.retrieve_online_documents_v2(**read_params) # type: ignore - ) + if request.api_version == 2: + response = await run_in_threadpool( + lambda: store.retrieve_online_documents_v2(**read_params) # type: ignore + ) + else: + response = await run_in_threadpool( + lambda: store.retrieve_online_documents(**read_params) # type: ignore + ) # Convert the Protobuf object to JSON and return it response_dict = await run_in_threadpool( diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 326daa589e8..cfad8178c05 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -2413,6 +2413,11 @@ def _retrieve_from_online_store_v2( output_len=output_len, ) + utils._populate_result_rows_from_columnar( + online_features_response=online_features_response, + data=entity_key_dict, + ) + return OnlineResponse(online_features_response) def serve( diff --git a/sdk/python/feast/infra/online_stores/remote.py b/sdk/python/feast/infra/online_stores/remote.py index ea09362299d..ec2b05759ba 100644 --- a/sdk/python/feast/infra/online_stores/remote.py +++ b/sdk/python/feast/infra/online_stores/remote.py @@ -31,6 +31,7 @@ feast_value_type_to_python_type, python_values_to_proto_values, ) +from feast.utils import _get_feature_view_vector_field_metadata from feast.value_type import ValueType logger = logging.getLogger(__name__) @@ -170,6 +171,235 @@ def online_read( logger.error(error_msg) raise RuntimeError(error_msg) + def retrieve_online_documents( + self, + config: RepoConfig, + table: FeatureView, + requested_features: Optional[List[str]], + embedding: Optional[List[float]], + top_k: int, + distance_metric: Optional[str] = "L2", + ) -> List[ + Tuple[ + Optional[datetime], + Optional[EntityKeyProto], + Optional[ValueProto], + Optional[ValueProto], + Optional[ValueProto], + ] + ]: + assert isinstance(config.online_store, RemoteOnlineStoreConfig) + config.online_store.__class__ = RemoteOnlineStoreConfig + + req_body = self._construct_online_documents_api_json_request( + table, requested_features, embedding, top_k, distance_metric + ) + response = get_remote_online_documents(config=config, req_body=req_body) + if response.status_code == 200: + logger.debug("Able to retrieve the online documents from feature server.") + response_json = json.loads(response.text) + event_ts: Optional[datetime] = self._get_event_ts(response_json) + + # Create feature name to index mapping for efficient lookup + feature_name_to_index = { + name: idx + for idx, name in enumerate(response_json["metadata"]["feature_names"]) + } + + vector_field_metadata = _get_feature_view_vector_field_metadata(table) + + # Process each result row + num_results = len(response_json["results"][0]["values"]) + result_tuples = [] + + for row_idx in range(num_results): + # Extract values using helper methods + feature_val = self._extract_requested_feature_value( + response_json, feature_name_to_index, requested_features, row_idx + ) + vector_value = self._extract_vector_field_value( + response_json, feature_name_to_index, vector_field_metadata, row_idx + ) + distance_val = self._extract_distance_value( + response_json, feature_name_to_index, "distance", row_idx + ) + entity_key_proto = self._construct_entity_key_from_response( + response_json, row_idx, feature_name_to_index, table + ) + + result_tuples.append( + ( + event_ts, + entity_key_proto, + feature_val, + vector_value, + distance_val, + ) + ) + + return result_tuples + else: + error_msg = f"Unable to retrieve the online documents using feature server API. Error_code={response.status_code}, error_message={response.text}" + logger.error(error_msg) + raise RuntimeError(error_msg) + + def retrieve_online_documents_v2( + self, + config: RepoConfig, + table: FeatureView, + requested_features: Optional[List[str]], + embedding: Optional[List[float]], + top_k: int, + distance_metric: Optional[str] = None, + query_string: Optional[str] = None, + ) -> List[ + Tuple[ + Optional[datetime], + Optional[EntityKeyProto], + Optional[Dict[str, ValueProto]], + ] + ]: + assert isinstance(config.online_store, RemoteOnlineStoreConfig) + config.online_store.__class__ = RemoteOnlineStoreConfig + + req_body = self._construct_online_documents_v2_api_json_request( + table, + requested_features, + embedding, + top_k, + distance_metric, + query_string, + api_version=2, + ) + response = get_remote_online_documents(config=config, req_body=req_body) + if response.status_code == 200: + logger.debug("Able to retrieve the online documents from feature server.") + response_json = json.loads(response.text) + event_ts: Optional[datetime] = self._get_event_ts(response_json) + + # Create feature name to index mapping for efficient lookup + feature_name_to_index = { + name: idx + for idx, name in enumerate(response_json["metadata"]["feature_names"]) + } + + # Process each result row + num_results = ( + len(response_json["results"][0]["values"]) + if response_json["results"] + else 0 + ) + result_tuples = [] + + for row_idx in range(num_results): + # Build feature values dictionary for requested features + feature_values_dict = {} + + if requested_features: + for feature_name in requested_features: + if feature_name in feature_name_to_index: + feature_idx = feature_name_to_index[feature_name] + if self._is_feature_present( + response_json, feature_idx, row_idx + ): + feature_values_dict[feature_name] = ( + self._extract_feature_value( + response_json, feature_idx, row_idx + ) + ) + else: + feature_values_dict[feature_name] = ValueProto() + + # Construct entity key proto using existing helper method + entity_key_proto = self._construct_entity_key_from_response( + response_json, row_idx, feature_name_to_index, table + ) + + result_tuples.append( + ( + event_ts, + entity_key_proto, + feature_values_dict if feature_values_dict else None, + ) + ) + + return result_tuples + else: + error_msg = f"Unable to retrieve the online documents using feature server API. Error_code={response.status_code}, error_message={response.text}" + logger.error(error_msg) + raise RuntimeError(error_msg) + + def _extract_requested_feature_value( + self, + response_json: dict, + feature_name_to_index: dict, + requested_features: Optional[List[str]], + row_idx: int, + ) -> Optional[ValueProto]: + """Extract the first available requested feature value.""" + if not requested_features: + return ValueProto() + + for feature_name in requested_features: + if feature_name in feature_name_to_index: + feature_idx = feature_name_to_index[feature_name] + if self._is_feature_present(response_json, feature_idx, row_idx): + return self._extract_feature_value( + response_json, feature_idx, row_idx + ) + + return ValueProto() + + def _extract_vector_field_value( + self, + response_json: dict, + feature_name_to_index: dict, + vector_field_metadata, + row_idx: int, + ) -> Optional[ValueProto]: + """Extract vector field value from response.""" + if ( + not vector_field_metadata + or vector_field_metadata.name not in feature_name_to_index + ): + return ValueProto() + + vector_feature_idx = feature_name_to_index[vector_field_metadata.name] + if self._is_feature_present(response_json, vector_feature_idx, row_idx): + return self._extract_feature_value( + response_json, vector_feature_idx, row_idx + ) + + return ValueProto() + + def _extract_distance_value( + self, + response_json: dict, + feature_name_to_index: dict, + distance_feature_name: str, + row_idx: int, + ) -> Optional[ValueProto]: + """Extract distance/score value from response.""" + if not distance_feature_name: + return ValueProto() + + distance_feature_idx = feature_name_to_index[distance_feature_name] + if self._is_feature_present(response_json, distance_feature_idx, row_idx): + distance_value = response_json["results"][distance_feature_idx]["values"][ + row_idx + ] + distance_val = ValueProto() + distance_val.float_val = float(distance_value) + return distance_val + + return ValueProto() + + def _is_feature_present( + self, response_json: dict, feature_idx: int, row_idx: int + ) -> bool: + """Check if a feature is present in the response.""" + return response_json["results"][feature_idx]["statuses"][row_idx] == "PRESENT" + def _construct_online_read_api_json_request( self, entity_keys: List[EntityKeyProto], @@ -197,12 +427,104 @@ def _construct_online_read_api_json_request( ) return req_body + def _construct_online_documents_api_json_request( + self, + table: FeatureView, + requested_features: Optional[List[str]] = None, + embedding: Optional[List[float]] = None, + top_k: Optional[int] = None, + distance_metric: Optional[str] = "L2", + ) -> str: + api_requested_features = [] + if requested_features is not None: + for requested_feature in requested_features: + api_requested_features.append(f"{table.name}:{requested_feature}") + + req_body = json.dumps( + { + "features": api_requested_features, + "query": embedding, + "top_k": top_k, + "distance_metric": distance_metric, + } + ) + return req_body + + def _construct_online_documents_v2_api_json_request( + self, + table: FeatureView, + requested_features: Optional[List[str]], + embedding: Optional[List[float]], + top_k: int, + distance_metric: Optional[str] = None, + query_string: Optional[str] = None, + api_version: Optional[int] = 2, + ) -> str: + api_requested_features = [] + if requested_features is not None: + for requested_feature in requested_features: + api_requested_features.append(f"{table.name}:{requested_feature}") + + req_body = json.dumps( + { + "features": api_requested_features, + "query": embedding, + "top_k": top_k, + "distance_metric": distance_metric, + "query_string": query_string, + "api_version": api_version, + } + ) + return req_body + def _get_event_ts(self, response_json) -> datetime: event_ts = "" if len(response_json["results"]) > 1: event_ts = response_json["results"][1]["event_timestamps"][0] return datetime.fromisoformat(event_ts.replace("Z", "+00:00")) + def _construct_entity_key_from_response( + self, + response_json: dict, + row_idx: int, + feature_name_to_index: dict, + table: FeatureView, + ) -> Optional[EntityKeyProto]: + """Construct EntityKeyProto from response data.""" + # Use the feature view's join_keys to identify entity fields + entity_fields = [ + join_key + for join_key in table.join_keys + if join_key in feature_name_to_index + ] + + if not entity_fields: + return None + + entity_key_proto = EntityKeyProto() + entity_key_proto.join_keys.extend(entity_fields) + + for entity_field in entity_fields: + if entity_field in feature_name_to_index: + feature_idx = feature_name_to_index[entity_field] + if self._is_feature_present(response_json, feature_idx, row_idx): + entity_value = self._extract_feature_value( + response_json, feature_idx, row_idx + ) + entity_key_proto.entity_values.append(entity_value) + + return entity_key_proto if entity_key_proto.entity_values else None + + def _extract_feature_value( + self, response_json: dict, feature_idx: int, row_idx: int + ) -> ValueProto: + """Extract and convert a feature value to ValueProto.""" + raw_value = response_json["results"][feature_idx]["values"][row_idx] + if raw_value is None: + return ValueProto() + proto_values = python_values_to_proto_values([raw_value]) + return proto_values[0] + def update( self, config: RepoConfig, @@ -239,6 +561,22 @@ def get_remote_online_features( ) +@rest_error_handling_decorator +def get_remote_online_documents( + session: requests.Session, config: RepoConfig, req_body: str +) -> requests.Response: + if config.online_store.cert: + return session.post( + f"{config.online_store.path}/retrieve-online-documents", + data=req_body, + verify=config.online_store.cert, + ) + else: + return session.post( + f"{config.online_store.path}/retrieve-online-documents", data=req_body + ) + + @rest_error_handling_decorator def post_remote_online_write( session: requests.Session, config: RepoConfig, req_body: dict diff --git a/sdk/python/feast/utils.py b/sdk/python/feast/utils.py index c63dad6a6ab..1f629f61f21 100644 --- a/sdk/python/feast/utils.py +++ b/sdk/python/feast/utils.py @@ -1048,7 +1048,7 @@ def _list_feature_views( def _get_feature_views_to_use( registry: "BaseRegistry", project, - features: Optional[Union[List[str], "FeatureService"]], + features: Union[List[str], "FeatureService"], allow_cache=False, hide_dummy_entity: bool = True, ) -> Tuple[List["FeatureView"], List["OnDemandFeatureView"]]: diff --git a/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py b/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py new file mode 100644 index 00000000000..1c074a40d40 --- /dev/null +++ b/sdk/python/tests/unit/infra/online_store/test_remote_online_store.py @@ -0,0 +1,353 @@ +import json +from datetime import datetime, timedelta +from unittest.mock import Mock, patch + +import pytest + +from feast import Entity, FeatureView, Field, FileSource, RepoConfig +from feast.infra.online_stores.remote import RemoteOnlineStore, RemoteOnlineStoreConfig +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.types import Float32, Int64, String +from feast.value_type import ValueType + + +class TestRemoteOnlineStoreRetrieveDocuments: + """Test suite for retrieve_online_documents and retrieve_online_documents_v2 methods.""" + + @pytest.fixture + def remote_store(self): + """Create a RemoteOnlineStore instance for testing.""" + return RemoteOnlineStore() + + @pytest.fixture + def config(self): + """Create a RepoConfig with RemoteOnlineStoreConfig.""" + return RepoConfig( + project="test_project", + online_store=RemoteOnlineStoreConfig( + type="remote", path="http://localhost:6566" + ), + registry="dummy_registry", + ) + + @pytest.fixture + def config_with_cert(self): + """Create a RepoConfig with RemoteOnlineStoreConfig including TLS cert.""" + return RepoConfig( + project="test_project", + online_store=RemoteOnlineStoreConfig( + type="remote", path="http://localhost:6566", cert="/path/to/cert.pem" + ), + registry="dummy_registry", + ) + + @pytest.fixture + def feature_view(self): + """Create a test FeatureView.""" + entity = Entity( + name="user_id", description="User ID", value_type=ValueType.INT64 + ) + source = FileSource(path="test.parquet", timestamp_field="event_timestamp") + return FeatureView( + name="test_feature_view", + entities=[entity], + ttl=timedelta(days=1), + schema=[ + Field(name="user_id", dtype=Int64), # Entity field + Field(name="feature1", dtype=String), + Field(name="embedding", dtype=Float32), + ], + source=source, + ) + + @pytest.fixture + def mock_successful_response(self): + """Create a mock successful HTTP response for documents retrieval.""" + return { + "metadata": { + "feature_names": ["feature1", "embedding", "distance", "user_id"] + }, + "results": [ + { + "values": ["test_value_1", "test_value_2"], + "statuses": ["PRESENT", "PRESENT"], + }, # feature1 + { + "values": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + "statuses": ["PRESENT", "PRESENT"], + "event_timestamps": [ + "2023-01-01T00:00:00Z", + "2023-01-01T01:00:00Z", + ], + }, # embedding + { + "values": [0.85, 0.92], + "statuses": ["PRESENT", "PRESENT"], + }, # distance + {"values": [123, 456], "statuses": ["PRESENT", "PRESENT"]}, # user_id + ], + } + + @pytest.fixture + def mock_successful_response_v2(self): + """Create a mock successful HTTP response for documents retrieval v2.""" + return { + "metadata": {"feature_names": ["user_id", "feature1"]}, + "results": [ + {"values": [123, 456], "statuses": ["PRESENT", "PRESENT"]}, # user_id + { + "values": ["test_value_1", "test_value_2"], + "statuses": ["PRESENT", "PRESENT"], + "event_timestamps": [ + "2023-01-01T00:00:00Z", + "2023-01-01T01:00:00Z", + ], + }, # feature1 + ], + } + + @patch("feast.infra.online_stores.remote.get_remote_online_documents") + def test_retrieve_online_documents_success( + self, + mock_get_remote_online_documents, + remote_store, + config, + feature_view, + mock_successful_response, + ): + """Test successful retrieve_online_documents call.""" + # Setup mock response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(mock_successful_response) + mock_get_remote_online_documents.return_value = mock_response + + # Call the method + result = remote_store.retrieve_online_documents( + config=config, + table=feature_view, + requested_features=["feature1"], + embedding=[0.1, 0.2, 0.3], + top_k=2, + distance_metric="L2", + ) + + # Verify the call was made correctly + mock_get_remote_online_documents.assert_called_once() + call_args = mock_get_remote_online_documents.call_args + assert call_args[1]["config"] == config + + # Parse the request body to verify it's correct + req_body = json.loads(call_args[1]["req_body"]) + assert req_body["features"] == ["test_feature_view:feature1"] + assert req_body["query"] == [0.1, 0.2, 0.3] + assert req_body["top_k"] == 2 + assert req_body["distance_metric"] == "L2" + + # Verify the result + assert len(result) == 2 + event_ts, entity_key_proto, feature_val, vector_value, distance_val = result[0] + + # Check event timestamp + assert isinstance(event_ts, datetime) + + # Check that we got ValueProto objects + assert isinstance(feature_val, ValueProto) + assert isinstance(vector_value, ValueProto) + assert isinstance(distance_val, ValueProto) + + @patch("feast.infra.online_stores.remote.get_remote_online_documents") + def test_retrieve_online_documents_v2_success( + self, + mock_get_remote_online_documents, + remote_store, + config, + feature_view, + mock_successful_response_v2, + ): + """Test successful retrieve_online_documents_v2 call.""" + # Setup mock response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(mock_successful_response_v2) + mock_get_remote_online_documents.return_value = mock_response + + # Call the method + result = remote_store.retrieve_online_documents_v2( + config=config, + table=feature_view, + requested_features=["feature1"], + embedding=[0.1, 0.2, 0.3], + top_k=2, + distance_metric="cosine", + query_string="test query", + ) + + # Verify the call was made correctly + mock_get_remote_online_documents.assert_called_once() + call_args = mock_get_remote_online_documents.call_args + assert call_args[1]["config"] == config + + # Parse the request body to verify it's correct + req_body = json.loads(call_args[1]["req_body"]) + assert req_body["features"] == ["test_feature_view:feature1"] + assert req_body["query"] == [0.1, 0.2, 0.3] + assert req_body["top_k"] == 2 + assert req_body["distance_metric"] == "cosine" + assert req_body["query_string"] == "test query" + assert req_body["api_version"] == 2 + + # Verify the result + assert len(result) == 2 + event_ts, entity_key_proto, feature_values_dict = result[0] + + # Check event timestamp + assert isinstance(event_ts, datetime) + + # Check entity key proto + assert isinstance(entity_key_proto, EntityKeyProto) + + # Check feature values dictionary + assert isinstance(feature_values_dict, dict) + assert "feature1" in feature_values_dict + assert isinstance(feature_values_dict["feature1"], ValueProto) + + @patch("feast.infra.online_stores.remote.get_remote_online_documents") + def test_retrieve_online_documents_with_cert( + self, + mock_get_remote_online_documents, + remote_store, + config_with_cert, + feature_view, + mock_successful_response, + ): + """Test retrieve_online_documents with TLS certificate.""" + # Setup mock response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(mock_successful_response) + mock_get_remote_online_documents.return_value = mock_response + + # Call the method + result = remote_store.retrieve_online_documents( + config=config_with_cert, + table=feature_view, + requested_features=["feature1"], + embedding=[0.1, 0.2, 0.3], + top_k=1, + ) + + # Verify the call was made + mock_get_remote_online_documents.assert_called_once() + assert len(result) == 2 + + @patch("feast.infra.online_stores.remote.get_remote_online_documents") + def test_retrieve_online_documents_error_response( + self, mock_get_remote_online_documents, remote_store, config, feature_view + ): + """Test retrieve_online_documents with error response.""" + # Setup mock error response + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + mock_get_remote_online_documents.return_value = mock_response + + # Call the method and expect RuntimeError + with pytest.raises( + RuntimeError, + match="Unable to retrieve the online documents using feature server API", + ): + remote_store.retrieve_online_documents( + config=config, + table=feature_view, + requested_features=["feature1"], + embedding=[0.1, 0.2, 0.3], + top_k=1, + ) + + @patch("feast.infra.online_stores.remote.get_remote_online_documents") + def test_retrieve_online_documents_v2_error_response( + self, mock_get_remote_online_documents, remote_store, config, feature_view + ): + """Test retrieve_online_documents_v2 with error response.""" + # Setup mock error response + mock_response = Mock() + mock_response.status_code = 404 + mock_response.text = "Not Found" + mock_get_remote_online_documents.return_value = mock_response + + # Call the method and expect RuntimeError + with pytest.raises( + RuntimeError, + match="Unable to retrieve the online documents using feature server API", + ): + remote_store.retrieve_online_documents_v2( + config=config, + table=feature_view, + requested_features=["feature1"], + embedding=[0.1, 0.2, 0.3], + top_k=1, + ) + + def test_construct_online_documents_api_json_request( + self, remote_store, feature_view + ): + """Test _construct_online_documents_api_json_request method.""" + result = remote_store._construct_online_documents_api_json_request( + table=feature_view, + requested_features=["feature1", "feature2"], + embedding=[0.1, 0.2, 0.3], + top_k=5, + distance_metric="cosine", + ) + + parsed_result = json.loads(result) + assert parsed_result["features"] == [ + "test_feature_view:feature1", + "test_feature_view:feature2", + ] + assert parsed_result["query"] == [0.1, 0.2, 0.3] + assert parsed_result["top_k"] == 5 + assert parsed_result["distance_metric"] == "cosine" + + def test_construct_online_documents_v2_api_json_request( + self, remote_store, feature_view + ): + """Test _construct_online_documents_v2_api_json_request method.""" + result = remote_store._construct_online_documents_v2_api_json_request( + table=feature_view, + requested_features=["feature1"], + embedding=[0.1, 0.2], + top_k=3, + distance_metric="L2", + query_string="test query", + api_version=2, + ) + + parsed_result = json.loads(result) + assert parsed_result["features"] == ["test_feature_view:feature1"] + assert parsed_result["query"] == [0.1, 0.2] + assert parsed_result["top_k"] == 3 + assert parsed_result["distance_metric"] == "L2" + assert parsed_result["query_string"] == "test query" + assert parsed_result["api_version"] == 2 + + def test_extract_requested_feature_value(self, remote_store): + """Test _extract_requested_feature_value helper method.""" + response_json = { + "results": [{"values": ["test_value"], "statuses": ["PRESENT"]}] + } + feature_name_to_index = {"feature1": 0} + + result = remote_store._extract_requested_feature_value( + response_json, feature_name_to_index, ["feature1"], 0 + ) + assert isinstance(result, ValueProto) + + def test_is_feature_present(self, remote_store): + """Test _is_feature_present helper method.""" + response_json = {"results": [{"statuses": ["PRESENT", "NOT_FOUND"]}]} + + assert remote_store._is_feature_present(response_json, 0, 0) + assert not remote_store._is_feature_present(response_json, 0, 1)