Exemplo n.º 1
0
    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
        )
Exemplo n.º 2
0
    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"
        )
Exemplo n.º 3
0
    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
        )
Exemplo n.º 4
0
    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"
        )