Remove value_type from tests

Signed-off-by: Felix Wang <wangfelix98@gmail.com>
feast-dev · feast-ci-bot · May 18, 2022 · Apr 25, 2022 · Apr 26, 2022 · Apr 26, 2022
commit 10523c77e2a6384599803bb25699d64c90407bf6
@@ -11,15 +11,13 @@ def setup_feature_store():
     """Prepares the local environment for a FeatureStore docstring test."""
     from datetime import datetime, timedelta
 
-    from feast import Entity, FeatureStore, FeatureView, Field, FileSource, ValueType
+    from feast import Entity, FeatureStore, FeatureView, Field, FileSource
     from feast.repo_operations import init_repo
     from feast.types import Float32, Int64
 
     init_repo("feature_repo", "local")
     fs = FeatureStore(repo_path="feature_repo")
-    driver = Entity(
-        name="driver_id", value_type=ValueType.INT64, description="driver id",
-    )
+    driver = Entity(name="driver_id", description="driver id",)
     driver_hourly_stats = FileSource(
         path="feature_repo/data/driver_stats.parquet",
         timestamp_field="event_timestamp",

@@ -1,14 +1,6 @@
 from datetime import timedelta
 
-from feast import (
-    BigQuerySource,
-    Entity,
-    FeatureService,
-    FeatureView,
-    Field,
-    PushSource,
-    ValueType,
-)
+from feast import BigQuerySource, Entity, FeatureService, FeatureView, Field, PushSource
 from feast.types import Float32, Int64, String
 
 driver_locations_source = BigQuerySource(
@@ -46,14 +38,12 @@
 driver = Entity(
     name="driver",  # The name is derived from this argument, not object name.
     join_keys=["driver_id"],
-    value_type=ValueType.INT64,
     description="driver id",
 )
 
 customer = Entity(
     name="customer",  # The name is derived from this argument, not object name.
     join_keys=["customer_id"],
-    value_type=ValueType.STRING,
 )
 
 

@@ -1,6 +1,6 @@
 from datetime import timedelta
 
-from feast import Entity, FeatureView, Field, FileSource, ValueType
+from feast import Entity, FeatureView, Field, FileSource
 from feast.types import Float32, Int32, Int64
 
 driver_hourly_stats = FileSource(
@@ -9,7 +9,7 @@
     created_timestamp_column="created",
 )
 
-driver = Entity(name="driver_id", value_type=ValueType.INT64, description="driver id")
+driver = Entity(name="driver_id", description="driver id",)
 
 
 driver_hourly_stats_view = FeatureView(

@@ -1,17 +1,12 @@
 from datetime import timedelta
 
-from feast import Entity, FeatureView, FileSource, ValueType
+from feast import Entity, FeatureView, FileSource
 
 driver_hourly_stats = FileSource(
     path="driver_stats.parquet",  # this parquet is not real and will not be read
 )
 
-driver = Entity(
-    name="driver_id",
-    value_type=ValueType.INT64,
-    description="driver id",
-    join_keys=["driver"],
-)
+driver = Entity(name="driver_id", description="driver id", join_keys=["driver"],)
 
 driver_hourly_stats_view = FeatureView(
     name="driver_hourly_stats",  # Intentionally use the same FeatureView name

@@ -1,6 +1,6 @@
 from datetime import timedelta
 
-from feast import Entity, FeatureView, Field, FileSource, ValueType
+from feast import Entity, FeatureView, Field, FileSource
 from feast.types import Float32, Int64
 
 driver_hourly_stats = FileSource(
@@ -11,12 +11,7 @@
 
 
 # The join key here is deliberately different from the parquet file to test the failure path.
-driver = Entity(
-    name="driver_id",
-    value_type=ValueType.INT64,
-    description="driver id",
-    join_keys=["driver"],
-)
+driver = Entity(name="driver_id", description="driver id", join_keys=["driver"],)
 
 
 driver_hourly_stats_view = FeatureView(

@@ -19,7 +19,7 @@
 
 import pytest
 
-from feast import Entity, RepoConfig, ValueType
+from feast import Entity, RepoConfig
 from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
 
 
@@ -60,7 +60,6 @@ def test_usage_on(dummy_exporter, enabling_toggle):
         entity = Entity(
             name="driver_car_id",
             description="Car driver id",
-            value_type=ValueType.STRING,
             tags={"team": "matchmaking"},
         )
 
@@ -99,7 +98,6 @@ def test_usage_off(dummy_exporter, enabling_toggle):
         entity = Entity(
             name="driver_car_id",
             description="Car driver id",
-            value_type=ValueType.STRING,
             tags={"team": "matchmaking"},
         )
         test_feature_store.apply([entity])

@@ -1,22 +1,21 @@
-from feast import Entity, ValueType
+from feast import Entity
 
 
-def driver(value_type: ValueType = ValueType.INT64):
+def driver():
     return Entity(
         name="driver",  # The name is derived from this argument, not object name.
-        value_type=value_type,
         description="driver id",
         join_keys=["driver_id"],
     )
 
 
 def customer():
-    return Entity(name="customer_id", value_type=ValueType.INT64)
+    return Entity(name="customer_id")
 
 
 def location():
-    return Entity(name="location_id", value_type=ValueType.INT64)
+    return Entity(name="location_id")
 
 
 def item():
-    return Entity(name="item_id", value_type=ValueType.INT64)
+    return Entity(name="item_id")
@@ -28,11 +28,10 @@ def driver_feature_view(
     name="test_correctness",
     infer_features: bool = False,
     dtype: FeastType = Float32,
-    entities: Optional[List[str]] = None,
 ) -> FeatureView:
     return FeatureView(
         name=name,
-        entities=entities or [driver()],
+        entities=[driver()],
         schema=None if infer_features else [Field(name="value", dtype=dtype)],
         ttl=timedelta(days=5),
         source=data_source,

@@ -728,7 +728,7 @@ def test_historical_features_from_bigquery_sources_containing_backfills(environm
         created_timestamp_column="created",
     )
 
-    driver = Entity(name="driver", join_keys=["driver_id"], value_type=ValueType.INT64)
+    driver = Entity(name="driver", join_keys=["driver_id"])
     driver_fv = FeatureView(
         name="driver_stats",
         entities=[driver],

@@ -13,7 +13,7 @@
 import requests
 from botocore.exceptions import BotoCoreError
 
-from feast import Entity, FeatureService, FeatureView, Field, ValueType
+from feast import Entity, FeatureService, FeatureView, Field
 from feast.errors import (
     FeatureNameCollisionError,
     RequestDataNotFoundInEntityRowsException,
@@ -117,7 +117,7 @@ def test_write_to_online_store_event_check(environment):
     }
     dataframe_source = pd.DataFrame(data)
     with prep_file_source(df=dataframe_source, timestamp_field="ts_1") as file_source:
-        e = Entity(name="id", value_type=ValueType.STRING)
+        e = Entity(name="id")
 
         # Create Feature View
         fv1 = FeatureView(

@@ -27,10 +27,8 @@
 from feast.infra.offline_stores.file import FileOfflineStoreConfig
 from feast.infra.online_stores.dynamodb import DynamoDBOnlineStoreConfig
 from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
-from feast.protos.feast.types import Value_pb2 as ValueProto
 from feast.repo_config import RepoConfig
 from feast.types import Array, Bytes, Float64, Int64, String
-from feast.value_type import ValueType
 from tests.utils.data_source_utils import (
     prep_file_source,
     simple_bq_source_using_query_arg,
@@ -93,10 +91,7 @@ def feature_store_with_s3_registry():
 )
 def test_apply_entity_success(test_feature_store):
     entity = Entity(
-        name="driver_car_id",
-        description="Car driver id",
-        value_type=ValueType.STRING,
-        tags={"team": "matchmaking"},
+        name="driver_car_id", description="Car driver id", tags={"team": "matchmaking"},
     )
 
     # Register Entity
@@ -108,7 +103,6 @@ def test_apply_entity_success(test_feature_store):
     assert (
         len(entities) == 1
         and entity.name == "driver_car_id"
-        and entity.value_type == ValueType(ValueProto.ValueType.STRING)
         and entity.description == "Car driver id"
         and "team" in entity.tags
         and entity.tags["team"] == "matchmaking"
@@ -127,10 +121,7 @@ def test_apply_entity_success(test_feature_store):
 )
 def test_apply_entity_integration(test_feature_store):
     entity = Entity(
-        name="driver_car_id",
-        description="Car driver id",
-        value_type=ValueType.STRING,
-        tags={"team": "matchmaking"},
+        name="driver_car_id", description="Car driver id", tags={"team": "matchmaking"},
     )
 
     # Register Entity
@@ -142,7 +133,6 @@ def test_apply_entity_integration(test_feature_store):
     assert (
         len(entities) == 1
         and entity.name == "driver_car_id"
-        and entity.value_type == ValueType(ValueProto.ValueType.STRING)
         and entity.description == "Car driver id"
         and "team" in entity.tags
         and entity.tags["team"] == "matchmaking"
@@ -151,7 +141,6 @@ def test_apply_entity_integration(test_feature_store):
     entity = test_feature_store.get_entity("driver_car_id")
     assert (
         entity.name == "driver_car_id"
-        and entity.value_type == ValueType(ValueProto.ValueType.STRING)
         and entity.description == "Car driver id"
         and "team" in entity.tags
         and entity.tags["team"] == "matchmaking"
@@ -219,9 +208,7 @@ def test_apply_feature_view_success(test_feature_store):
 @pytest.mark.parametrize("dataframe_source", [lazy_fixture("simple_dataset_1")])
 def test_feature_view_inference_success(test_feature_store, dataframe_source):
     with prep_file_source(df=dataframe_source, timestamp_field="ts_1") as file_source:
-        entity = Entity(
-            name="id", join_keys=["id_join_key"], value_type=ValueType.INT64
-        )
+        entity = Entity(name="id", join_keys=["id_join_key"])
 
         fv1 = FeatureView(
             name="fv1",
@@ -368,13 +355,9 @@ def test_apply_object_and_read(test_feature_store):
         created_timestamp_column="timestamp",
     )
 
-    e1 = Entity(
-        name="fs1_my_entity_1", value_type=ValueType.STRING, description="something"
-    )
+    e1 = Entity(name="fs1_my_entity_1", description="something")
 
-    e2 = Entity(
-        name="fs1_my_entity_2", value_type=ValueType.STRING, description="something"
-    )
+    e2 = Entity(name="fs1_my_entity_2", description="something")
 
     fv1 = FeatureView(
         name="my_feature_view_1",
@@ -438,7 +421,7 @@ def test_apply_remote_repo():
 def test_reapply_feature_view_success(test_feature_store, dataframe_source):
     with prep_file_source(df=dataframe_source, timestamp_field="ts_1") as file_source:
 
-        e = Entity(name="id", join_keys=["id_join_key"], value_type=ValueType.STRING)
+        e = Entity(name="id", join_keys=["id_join_key"])
 
         # Create Feature View
         fv1 = FeatureView(