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" )
def test_stop_ingest_job(self, mock_jobcontroller_client, mocker): mocker.patch.object( mock_jobcontroller_client, "_jobcontroller_service_stub", return_value=Core.JobControllerServiceStub( grpc.insecure_channel("")), ) ingest_job = IngestJob( job_proto=IngestJobProto( id="kafka-to-redis", external_id="job#2222", status=IngestionJobStatus.RUNNING, ), core_stub=mock_jobcontroller_client._jobcontroller_service_stub, ) mock_jobcontroller_client.stop_ingest_job(ingest_job) assert (mock_jobcontroller_client._jobcontroller_service_stub. StopIngestionJob.called)