From c81dc2795db3e3853b01a41eb1147f5ac4150839 Mon Sep 17 00:00:00 2001
From: Judah Rand <17158624+judahrand@users.noreply.github.com>
Date: Tue, 25 Jan 2022 10:24:14 +0000
Subject: [PATCH] Add simple unittest for `_get_unique_entities`

Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com>
---
 .../tests/unit/test_unit_feature_store.py     | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 sdk/python/tests/unit/test_unit_feature_store.py

diff --git a/sdk/python/tests/unit/test_unit_feature_store.py b/sdk/python/tests/unit/test_unit_feature_store.py
new file mode 100644
index 0000000000..6f9dd6acb0
--- /dev/null
+++ 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])