-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathtest_unit_feature_store.py
More file actions
131 lines (108 loc) · 3.84 KB
/
test_unit_feature_store.py
File metadata and controls
131 lines (108 loc) · 3.84 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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from dataclasses import dataclass
from typing import Dict, List
import pytest
from feast import utils
from feast.protos.feast.types.Value_pb2 import Value
@dataclass
class MockFeatureViewProjection:
join_key_map: Dict[str, str]
@dataclass
class MockFeatureView:
name: str
entities: List[str]
projection: MockFeatureViewProjection
def test_get_unique_entities_success():
entity_values = {
"entity_1": [Value(int64_val=1), Value(int64_val=2), Value(int64_val=1)],
"entity_2": [
Value(string_val="1"),
Value(string_val="2"),
Value(string_val="1"),
],
"entity_3": [Value(int64_val=8), Value(int64_val=9), Value(int64_val=10)],
}
entity_name_to_join_key_map = {"entity_1": "entity_1", "entity_2": "entity_2"}
fv = MockFeatureView(
name="fv_1",
entities=["entity_1", "entity_2"],
projection=MockFeatureViewProjection(join_key_map={}),
)
unique_entities, indexes, output_len = utils._get_unique_entities(
table=fv,
join_key_values=entity_values,
entity_name_to_join_key_map=entity_name_to_join_key_map,
)
expected_entities = (
{"entity_1": Value(int64_val=1), "entity_2": Value(string_val="1")},
{"entity_1": Value(int64_val=2), "entity_2": Value(string_val="2")},
)
expected_indexes = ([0, 2], [1])
assert unique_entities == expected_entities
assert indexes == expected_indexes
assert output_len == 3
def test_get_unique_entities_missing_join_key_success():
"""
Tests that _get_unique_entities raises a KeyError when a required join key is missing.
"""
# Here, we omit the required key for "entity_1"
entity_values = {
"entity_2": [
Value(string_val="1"),
Value(string_val="2"),
Value(string_val="1"),
],
}
entity_name_to_join_key_map = {"entity_1": "entity_1", "entity_2": "entity_2"}
fv = MockFeatureView(
name="fv_1",
entities=["entity_1", "entity_2"],
projection=MockFeatureViewProjection(join_key_map={}),
)
unique_entities, indexes, output_len = utils._get_unique_entities(
table=fv,
join_key_values=entity_values,
entity_name_to_join_key_map=entity_name_to_join_key_map,
)
expected_entities = (
{"entity_2": Value(string_val="1")},
{"entity_2": Value(string_val="2")},
)
expected_indexes = ([0, 2], [1])
assert unique_entities == expected_entities
assert indexes == expected_indexes
assert output_len == 3
# We're not say anything about the entity_1 missing from the unique_entities list
assert "entity_1" not in [entity.keys() for entity in unique_entities]
def test_get_unique_entities_missing_all_join_keys_error():
"""
Tests that _get_unique_entities raises a KeyError when all required join keys are missing.
"""
entity_values_not_in_feature_view = {
"entity_3": [Value(string_val="3")],
}
entity_name_to_join_key_map = {
"entity_1": "entity_1",
"entity_2": "entity_2",
"entity_3": "entity_3",
}
fv = MockFeatureView(
name="fv_1",
entities=["entity_1", "entity_2"],
projection=MockFeatureViewProjection(join_key_map={}),
)
with pytest.raises(KeyError) as excinfo:
utils._get_unique_entities(
table=fv,
join_key_values=entity_values_not_in_feature_view,
entity_name_to_join_key_map=entity_name_to_join_key_map,
)
error_message = str(excinfo.value)
assert (
"Missing join key values for keys: ['entity_1', 'entity_2', 'entity_3']"
in error_message
)
assert (
"No values provided for keys: ['entity_1', 'entity_2', 'entity_3']"
in error_message
)
assert "Provided join_key_values: ['entity_3']" in error_message