def test_construct(
self,
dataframe,
feature_set_dataframe,
key_id,
timestamp_c,
feature_add,
feature_divide,
):
spark_client = Mock()
# arrange
feature_set = FeatureSet(
"name",
"entity",
"description",
[key_id],
timestamp_c,
[feature_add, feature_divide],
)
# act
result_df = feature_set.construct(dataframe, spark_client)
result_columns = result_df.columns
# assert
assert (result_columns == key_id.get_output_columns() +
timestamp_c.get_output_columns() +
feature_add.get_output_columns() +
feature_divide.get_output_columns())
assert_dataframe_equality(result_df, feature_set_dataframe)
assert result_df.is_cached
def test_columns(self, key_id, timestamp_c, feature_add, feature_divide):
# arrange
name = "name"
entity = "entity"
description = "description"
# act
fs = FeatureSet(
name,
entity,
description,
[key_id],
timestamp_c,
[feature_add, feature_divide],
)
out_columns = fs.columns
# assert
assert (
out_columns
== key_id.get_output_columns()
+ timestamp_c.get_output_columns()
+ feature_add.get_output_columns()
+ feature_divide.get_output_columns()
)