Skip to content

Commit

Permalink
Add simple unittest for _get_unique_entities
Browse files Browse the repository at this point in the history
Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com>
  • Loading branch information
judahrand committed Jan 25, 2022
1 parent 39a5812 commit c81dc27
Showing 1 changed file with 50 additions and 0 deletions.
50 changes: 50 additions & 0 deletions sdk/python/tests/unit/test_unit_feature_store.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
from dataclasses import dataclass
from typing import Dict, List

from feast import FeatureStore
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():
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 = FeatureStore._get_unique_entities(
FeatureStore,
table=fv,
join_key_values=entity_values,
entity_name_to_join_key_map=entity_name_to_join_key_map,
)

assert unique_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")},
)
assert indexes == ([0, 2], [1])

0 comments on commit c81dc27

Please sign in to comment.