add test to ensure only calling async when supported for the push mode

robhowley · robhowley · commit 3b79f1fd4dc6 · 2024-10-18T17:22:40.000-04:00
diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -213,6 +213,7 @@ async def push(body=Depends(get_body)):
             and to in [PushMode.ONLINE, PushMode.ONLINE_AND_OFFLINE]
         )
         if should_push_async:
+            print("im in async?")
             await store.push_async(**push_params)
         else:
             store.push(**push_params)
diff --git a/sdk/python/tests/foo_provider.py b/sdk/python/tests/foo_provider.py
@@ -22,6 +22,10 @@
 from feast.infra.offline_stores.offline_store import RetrievalJob
 from feast.infra.provider import Provider
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.infra.supported_async_methods import (
+    ProviderAsyncMethods,
+    SupportedAsyncMethods,
+)
 from feast.online_response import OnlineResponse
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import RepeatedValue
@@ -30,6 +34,20 @@
 
 
 class FooProvider(Provider):
+    @staticmethod
+    def with_async_support(online_read=False, online_write=False):
+        class _FooProvider(FooProvider):
+            @property
+            def async_supported(self):
+                return ProviderAsyncMethods(
+                    online=SupportedAsyncMethods(
+                        read=online_read,
+                        write=online_write,
+                    )
+                )
+
+        return _FooProvider(None)
+
     def __init__(self, config: RepoConfig):
         pass
 
diff --git a/sdk/python/tests/unit/test_feature_server.py b/sdk/python/tests/unit/test_feature_server.py
@@ -0,0 +1,48 @@
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi.testclient import TestClient
+
+from feast import FeatureStore
+from feast.data_source import PushMode
+from feast.feature_server import get_app
+from feast.utils import _utc_now
+from tests.foo_provider import FooProvider
+
+
+@pytest.mark.parametrize(
+    "online_write,push_mode,async_count",
+    [
+        (True, PushMode.ONLINE_AND_OFFLINE, 1),
+        (True, PushMode.OFFLINE, 0),
+        (True, PushMode.ONLINE, 1),
+        (False, PushMode.ONLINE_AND_OFFLINE, 0),
+        (False, PushMode.OFFLINE, 0),
+        (False, PushMode.ONLINE, 0),
+    ],
+)
+def test_push_online_async_supported(online_write, push_mode, async_count, environment):
+    push_payload = json.dumps(
+        {
+            "push_source_name": "location_stats_push_source",
+            "df": {
+                "location_id": [1],
+                "temperature": [100],
+                "event_timestamp": [str(_utc_now())],
+                "created": [str(_utc_now())],
+            },
+            "to": push_mode.name.lower(),
+        }
+    )
+
+    provider = FooProvider.with_async_support(online_write=online_write)
+    print(provider.async_supported.online.write)
+    with patch.object(FeatureStore, "_get_provider", return_value=provider):
+        fs = environment.feature_store
+        fs.push = MagicMock()
+        fs.push_async = AsyncMock()
+        client = TestClient(get_app(fs))
+        client.post("/push", data=push_payload)
+        assert fs.push.call_count == 1 - async_count
+        assert fs.push_async.await_count == async_count

Original file line number	Diff line number	Diff line change
`@@ -213,6 +213,7 @@ async def push(body=Depends(get_body)):`
`213`	`213`	`and to in [PushMode.ONLINE, PushMode.ONLINE_AND_OFFLINE]`
`214`	`214`	`)`
`215`	`215`	`if should_push_async:`
	`216`	`+ print("im in async?")`
`216`	`217`	`await store.push_async(**push_params)`
`217`	`218`	`else:`
`218`	`219`	`store.push(**push_params)`