import pytest

from utils.data_source_utils import (

prep_file_source,

simple_bq_source_using_query_arg,

simple_bq_source_using_table_ref_arg,

)

from feast import Entity, ValueType

from feast.errors import RegistryInferenceFailure

from feast.feature_view import FeatureView

from feast.inference import (

infer_entity_value_type_from_feature_views,

update_data_sources_with_inferred_event_timestamp_col,

)

def test_infer_entity_value_type_from_feature_views(simple_dataset_1, simple_dataset_2):

with prep_file_source(

df=simple_dataset_1, event_timestamp_column="ts_1"

) as file_source, prep_file_source(

df=simple_dataset_2, event_timestamp_column="ts_1"

) as file_source_2:

fv1 = FeatureView(name="fv1", entities=["id"], input=file_source, ttl=None,)

fv2 = FeatureView(name="fv2", entities=["id"], input=file_source_2, ttl=None,)

actual_1 = infer_entity_value_type_from_feature_views(

[Entity(name="id")], [fv1]

)

actual_2 = infer_entity_value_type_from_feature_views(

[Entity(name="id")], [fv2]

)

assert actual_1 == [Entity(name="id", value_type=ValueType.INT64)]

assert actual_2 == [Entity(name="id", value_type=ValueType.STRING)]

with pytest.raises(RegistryInferenceFailure):

# two viable data types

infer_entity_value_type_from_feature_views([Entity(name="id")], [fv1, fv2])

@pytest.mark.integration

def test_infer_event_timestamp_column_for_data_source(simple_dataset_1):

df_with_two_viable_timestamp_cols = simple_dataset_1.copy(deep=True)

df_with_two_viable_timestamp_cols["ts_2"] = simple_dataset_1["ts_1"]

with prep_file_source(df=simple_dataset_1) as file_source:

data_sources = [

file_source,

simple_bq_source_using_table_ref_arg(simple_dataset_1),

simple_bq_source_using_query_arg(simple_dataset_1),

]

update_data_sources_with_inferred_event_timestamp_col(data_sources)

actual_event_timestamp_cols = [

source.event_timestamp_column for source in data_sources

]

assert actual_event_timestamp_cols == ["ts_1", "ts_1", "ts_1"]

with prep_file_source(df=df_with_two_viable_timestamp_cols) as file_source:

with pytest.raises(RegistryInferenceFailure):

# two viable event_timestamp_columns

update_data_sources_with_inferred_event_timestamp_col([file_source])