Merge pull request feast-dev#6 from redhatHameed/update_client

added unit test for offline store remote client

Author: dmartinol

examples/remote-offline-store/offline_server/feature_repo/data/online_store.db

@@ -35,15 +35,13 @@ class RemoteOfflineStoreConfig(FeastConfigBaseModel):
 class RemoteRetrievalJob(RetrievalJob):
     def __init__(
         self,
-        config: RepoConfig,
+        client: fl.FlightClient,
         feature_refs: List[str],
         entity_df: Union[pd.DataFrame, str],
         # TODO add missing parameters from the OfflineStore API
     ):
         # Initialize the client connection
-        self.client = fl.connect(
-            f"grpc://{config.offline_store.host}:{config.offline_store.port}"
-        )
+        self.client = client
         self.feature_refs = feature_refs
         self.entity_df = entity_df
 
@@ -108,8 +106,14 @@ def get_historical_features(
         assert isinstance(config.offline_store, RemoteOfflineStoreConfig)
 
         # TODO: extend RemoteRetrievalJob API with all method parameters
+
+        # Initialize the client connection
+        client = fl.connect(
+            f"grpc://{config.offline_store.host}:{config.offline_store.port}"
+        )
+
         return RemoteRetrievalJob(
-            config=config, feature_refs=feature_refs, entity_df=entity_df
+            client=client, feature_refs=feature_refs, entity_df=entity_df
         )
 
     @staticmethod

sdk/python/feast/infra/offline_stores/remote.py

@@ -30,6 +30,10 @@
     RedshiftOfflineStoreConfig,
     RedshiftRetrievalJob,
 )
+from feast.infra.offline_stores.remote import (
+    RemoteOfflineStoreConfig,
+    RemoteRetrievalJob,
+)
 from feast.infra.offline_stores.snowflake import (
     SnowflakeOfflineStoreConfig,
     SnowflakeRetrievalJob,
@@ -105,6 +109,7 @@ def metadata(self) -> Optional[RetrievalMetadata]:
         PostgreSQLRetrievalJob,
         SparkRetrievalJob,
         TrinoRetrievalJob,
+        RemoteRetrievalJob,
     ]
 )
 def retrieval_job(request, environment):
@@ -206,6 +211,26 @@ def retrieval_job(request, environment):
             config=environment.test_repo_config,
             full_feature_names=False,
         )
+    elif request.param is RemoteRetrievalJob:
+        offline_store_config = RemoteOfflineStoreConfig(
+            type="remote",
+            host="localhost",
+            port=0,
+        )
+        environment.test_repo_config.offline_store = offline_store_config
+        return RemoteRetrievalJob(
+            client=MagicMock(),
+            feature_refs=[
+                "str:str",
+            ],
+            entity_df=pd.DataFrame.from_dict(
+                {
+                    "id": [1],
+                    "event_timestamp": ["datetime"],
+                    "val_to_add": [1],
+                }
+            ),
+        )
     else:
         return request.param()