updated tests

mavysavydav · mavysavydav · commit 40df252912e2 · 2021-11-18T00:06:07.000-08:00
Signed-off-by: David Y Liu &lt;davidyliuliu@gmail.com&gt;
diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py
@@ -99,7 +99,7 @@ def simple_dataset_1() -> pd.DataFrame:
     now = datetime.utcnow()
     ts = pd.Timestamp(now).round("ms")
     data = {
-        "id": [1, 2, 1, 3, 3],
+        "id_join_key": [1, 2, 1, 3, 3],
         "float_col": [0.1, 0.2, 0.3, 4, 5],
         "int64_col": [1, 2, 3, 4, 5],
         "string_col": ["a", "b", "c", "d", "e"],
@@ -119,7 +119,7 @@ def simple_dataset_2() -> pd.DataFrame:
     now = datetime.utcnow()
     ts = pd.Timestamp(now).round("ms")
     data = {
-        "id": ["a", "b", "c", "d", "e"],
+        "id_join_key": ["a", "b", "c", "d", "e"],
         "float_col": [0.1, 0.2, 0.3, 4, 5],
         "int64_col": [1, 2, 3, 4, 5],
         "string_col": ["a", "b", "c", "d", "e"],
diff --git a/sdk/python/tests/integration/registration/test_feature_store.py b/sdk/python/tests/integration/registration/test_feature_store.py
@@ -218,6 +218,8 @@ def test_feature_view_inference_success(test_feature_store, dataframe_source):
     with prep_file_source(
         df=dataframe_source, event_timestamp_column="ts_1"
     ) as file_source:
+        entity = Entity(name="id", join_key="id_join_key", value_type=ValueType.INT64)
+
         fv1 = FeatureView(
             name="fv1",
             entities=["id"],
@@ -245,7 +247,7 @@ def test_feature_view_inference_success(test_feature_store, dataframe_source):
             tags={},
         )
 
-        test_feature_store.apply([fv1, fv2, fv3])  # Register Feature Views
+        test_feature_store.apply([entity, fv1, fv2, fv3])  # Register Feature Views
         feature_view_1 = test_feature_store.list_feature_views()[0]
         feature_view_2 = test_feature_store.list_feature_views()[1]
         feature_view_3 = test_feature_store.list_feature_views()[2]
@@ -433,7 +435,7 @@ def test_reapply_feature_view_success(test_feature_store, dataframe_source):
         df=dataframe_source, event_timestamp_column="ts_1"
     ) as file_source:
 
-        e = Entity(name="id", value_type=ValueType.STRING)
+        e = Entity(name="id", join_key="id_join_key", value_type=ValueType.STRING)
 
         # Create Feature View
         fv1 = FeatureView(
diff --git a/sdk/python/tests/integration/registration/test_inference.py b/sdk/python/tests/integration/registration/test_inference.py
@@ -30,22 +30,26 @@ def test_update_entities_with_inferred_types_from_feature_views(
             name="fv2", entities=["id"], batch_source=file_source_2, ttl=None,
         )
 
-        actual_1 = Entity(name="id")
-        actual_2 = Entity(name="id")
+        actual_1 = Entity(name="id", join_key="id_join_key")
+        actual_2 = Entity(name="id", join_key="id_join_key")
 
         update_entities_with_inferred_types_from_feature_views(
             [actual_1], [fv1], RepoConfig(provider="local", project="test")
         )
         update_entities_with_inferred_types_from_feature_views(
             [actual_2], [fv2], RepoConfig(provider="local", project="test")
         )
-        assert actual_1 == Entity(name="id", value_type=ValueType.INT64)
-        assert actual_2 == Entity(name="id", value_type=ValueType.STRING)
+        assert actual_1 == Entity(
+            name="id", join_key="id_join_key", value_type=ValueType.INT64
+        )
+        assert actual_2 == Entity(
+            name="id", join_key="id_join_key", value_type=ValueType.STRING
+        )
 
         with pytest.raises(RegistryInferenceFailure):
             # two viable data types
             update_entities_with_inferred_types_from_feature_views(
-                [Entity(name="id")],
+                [Entity(name="id", join_key="id_join_key")],
                 [fv1, fv2],
                 RepoConfig(provider="local", project="test"),
             )
