-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathtest_feature_server.py
More file actions
47 lines (42 loc) · 1.52 KB
/
Copy pathtest_feature_server.py
File metadata and controls
47 lines (42 loc) · 1.52 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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)
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