def test_get_feature_set(self, mocked_client, mocker):
mocked_client._core_service_stub = Core.CoreServiceStub(
grpc.insecure_channel("")
)
from google.protobuf.duration_pb2 import Duration
mocker.patch.object(
mocked_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(
feature_set=FeatureSetProto(
spec=FeatureSetSpecProto(
name="my_feature_set",
max_age=Duration(seconds=3600),
labels={"key1": "val1", "key2": "val2"},
features=[
FeatureSpecProto(
name="my_feature_1",
value_type=ValueProto.ValueType.FLOAT,
),
FeatureSpecProto(
name="my_feature_2",
value_type=ValueProto.ValueType.FLOAT,
),
],
entities=[
EntitySpecProto(
name="my_entity_1",
value_type=ValueProto.ValueType.INT64,
)
],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
),
meta=FeatureSetMetaProto(),
)
),
)
mocked_client.set_project("my_project")
feature_set = mocked_client.get_feature_set("my_feature_set")
assert (
feature_set.name == "my_feature_set"
and "key1" in feature_set.labels
and feature_set.labels["key1"] == "val1"
and "key2" in feature_set.labels
and feature_set.labels["key2"] == "val2"
and feature_set.fields["my_feature_1"].name == "my_feature_1"
and feature_set.fields["my_feature_1"].dtype == ValueType.FLOAT
and feature_set.fields["my_entity_1"].name == "my_entity_1"
and feature_set.fields["my_entity_1"].dtype == ValueType.INT64
and len(feature_set.features) == 2
and len(feature_set.entities) == 1
)
def test_list_ingest_jobs(self, mocked_client, mocker):
mocker.patch.object(
mocked_client,
"_core_service_stub",
return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
)
feature_set_proto = FeatureSetProto(
spec=FeatureSetSpecProto(
project="test", name="driver", max_age=Duration(seconds=3600),
)
)
mocker.patch.object(
mocked_client._core_service_stub,
"ListIngestionJobs",
return_value=ListIngestionJobsResponse(
jobs=[
IngestJobProto(
id="kafka-to-redis",
external_id="job-2222",
status=IngestionJobStatus.RUNNING,
feature_sets=[feature_set_proto],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
store=Store(name="redis"),
)
]
),
)
# list ingestion jobs by target feature set reference
ingest_jobs = mocked_client.list_ingest_jobs(
feature_set_ref=FeatureSetRef.from_feature_set(
FeatureSet.from_proto(feature_set_proto)
)
)
assert len(ingest_jobs) >= 1
ingest_job = ingest_jobs[0]
assert (
ingest_job.status == IngestionJobStatus.RUNNING
and ingest_job.id == "kafka-to-redis"
and ingest_job.external_id == "job-2222"
and ingest_job.feature_sets[0].name == "driver"
and ingest_job.source.source_type == "Kafka"
)
def test_get_feature_set(self, mock_client, mocker):
mock_client._core_service_stub = Core.CoreServiceStub(grpc.insecure_channel(""))
from google.protobuf.duration_pb2 import Duration
mocker.patch.object(
mock_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(
feature_set=FeatureSetSpec(
name="my_feature_set",
version=2,
max_age=Duration(seconds=3600),
features=[
FeatureSpec(
name="my_feature_1", value_type=ValueProto.ValueType.FLOAT
),
FeatureSpec(
name="my_feature_2", value_type=ValueProto.ValueType.FLOAT
),
],
entities=[
EntitySpec(
name="my_entity_1", value_type=ValueProto.ValueType.INT64
)
],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
)
),
)
feature_set = mock_client.get_feature_set("my_feature_set", version=2)
assert (
feature_set.name == "my_feature_set"
and feature_set.version == 2
and feature_set.fields["my_feature_1"].name == "my_feature_1"
and feature_set.fields["my_feature_1"].dtype == ValueType.FLOAT
and feature_set.fields["my_entity_1"].name == "my_entity_1"
and feature_set.fields["my_entity_1"].dtype == ValueType.INT64
and len(feature_set.features) == 2
and len(feature_set.entities) == 1
)
def test_list_ingest_jobs(self, mock_jobcontroller_client, mocker):
mocker.patch.object(
mock_jobcontroller_client,
"_jobcontroller_service_stub",
return_value=Core.JobControllerServiceStub(grpc.insecure_channel("")),
)
feature_set_ref = FeatureSetRef(project="test", name="driver",)
mocker.patch.object(
mock_jobcontroller_client._jobcontroller_service_stub,
"ListIngestionJobs",
return_value=ListIngestionJobsResponse(
jobs=[
IngestJobProto(
id="kafka-to-redis",
external_id="job-2222",
status=IngestionJobStatus.RUNNING,
feature_set_references=[feature_set_ref.to_proto()],
source=Source(
type=SourceType.KAFKA,
kafka_source_config=KafkaSourceConfig(
bootstrap_servers="localhost:9092", topic="topic"
),
),
stores=[Store(name="redis")],
)
]
),
)
# list ingestion jobs by target feature set reference
ingest_jobs = mock_jobcontroller_client.list_ingest_jobs(
feature_set_ref=feature_set_ref
)
assert len(ingest_jobs) >= 1
ingest_job = ingest_jobs[0]
assert (
ingest_job.status == IngestionJobStatus.RUNNING
and ingest_job.id == "kafka-to-redis"
and ingest_job.external_id == "job-2222"
and ingest_job.feature_sets[0].name == "driver"
and ingest_job.source.source_type == "Kafka"
)