def test_feature_set_ingest_success(self, dataframe, client): driver_fs = FeatureSet("driver-feature-set") driver_fs.add(Feature(name="feature_1", dtype=ValueType.FLOAT)) driver_fs.add(Feature(name="feature_2", dtype=ValueType.STRING)) driver_fs.add(Feature(name="feature_3", dtype=ValueType.INT64)) driver_fs.add(Entity(name="entity_id", dtype=ValueType.INT64)) driver_fs.source = KafkaSource(topic="feature-topic", brokers="127.0.0.1") driver_fs._message_producer = MagicMock() driver_fs._message_producer.send = MagicMock() # Register with Feast core client.apply(driver_fs) # Ingest data into Feast driver_fs.ingest(dataframe=dataframe) # Make sure message producer is called driver_fs._message_producer.send.assert_called()
def test_feature_set_types_success(self, client, dataframe): all_types_fs = FeatureSet( name="all_types", entities=[Entity(name="user_id", dtype=ValueType.INT64)], features=[ Feature(name="float_feature", dtype=ValueType.FLOAT), Feature(name="int64_feature", dtype=ValueType.INT64), Feature(name="int32_feature", dtype=ValueType.INT32), Feature(name="string_feature", dtype=ValueType.STRING), Feature(name="bytes_feature", dtype=ValueType.BYTES), Feature(name="bool_feature", dtype=ValueType.BOOL), Feature(name="double_feature", dtype=ValueType.DOUBLE), Feature(name="float_list_feature", dtype=ValueType.FLOAT_LIST), Feature(name="int64_list_feature", dtype=ValueType.INT64_LIST), Feature(name="int32_list_feature", dtype=ValueType.INT32_LIST), Feature(name="string_list_feature", dtype=ValueType.STRING_LIST), Feature(name="bytes_list_feature", dtype=ValueType.BYTES_LIST), Feature(name="bool_list_feature", dtype=ValueType.BOOL_LIST), Feature(name="double_list_feature", dtype=ValueType.DOUBLE_LIST), ], max_age=Duration(seconds=3600), ) all_types_fs.source = KafkaSource(topic="feature-topic", brokers="127.0.0.1") all_types_fs._message_producer = MagicMock() all_types_fs._message_producer.send = MagicMock() # Register with Feast core client.apply(all_types_fs) # Ingest data into Feast all_types_fs.ingest(dataframe=dataframe) # Make sure message producer is called all_types_fs._message_producer.send.assert_called()
def test_feature_set_ingest_failure(self, client, dataframe, exception): with pytest.raises(exception): # Create feature set driver_fs = FeatureSet("driver-feature-set") driver_fs.source = KafkaSource(topic="feature-topic", brokers="fake.broker.com") driver_fs._message_producer = MagicMock() driver_fs._message_producer.send = MagicMock() # Update based on dataset driver_fs.update_from_dataset( dataframe, column_mapping={ "entity_id": Entity(name="entity", dtype=ValueType.INT64) }, ) # Register with Feast core client.apply(driver_fs) # Ingest data into Feast driver_fs.ingest(dataframe=dataframe)