Add simple unittest for _get_unique_entities

judahrand · judahrand · commit c81dc2795db3 · 2022-01-25T10:30:03.000Z
Signed-off-by: Judah Rand &lt;17158624+judahrand@users.noreply.github.com&gt;
diff --git a/sdk/python/tests/unit/test_unit_feature_store.py b/sdk/python/tests/unit/test_unit_feature_store.py
@@ -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])
