From 40df252912e210582845568192d56322316a0b94 Mon Sep 17 00:00:00 2001 From: David Y Liu Date: Wed, 17 Nov 2021 17:45:32 -0800 Subject: [PATCH] updated tests Signed-off-by: David Y Liu --- sdk/python/tests/conftest.py | 4 ++-- .../integration/registration/test_feature_store.py | 6 ++++-- .../integration/registration/test_inference.py | 14 +++++++++----- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py index 63bf4be629..b7731fbb4c 100644 --- 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 index b0cedfd642..d605865960 100644 --- 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 index f45a18ea55..ca0a640849 100644 --- a/sdk/python/tests/integration/registration/test_inference.py +++ b/sdk/python/tests/integration/registration/test_inference.py @@ -30,8 +30,8 @@ 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") @@ -39,13 +39,17 @@ def test_update_entities_with_inferred_types_from_feature_views( 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"), )