def test_restarting_failed_jobs(feature_table): """ If configured - restart failed jobs """ feast_client = FeastClient( job_service_pause_between_jobs=0, job_service_retry_failed_jobs=True, options={"whitelisted_projects": "default,ride"}, ) feast_client.list_projects = Mock(return_value=["default"]) feast_client.list_feature_tables = Mock() spark_client = Client(feast_client) spark_client.list_jobs = Mock() spark_client.start_stream_to_online_ingestion = Mock() spark_client.feature_store.list_feature_tables.return_value = [ feature_table ] spark_client.list_jobs.return_value = [] ensure_stream_ingestion_jobs(spark_client, all_projects=True) spark_client.list_jobs.assert_called_once_with(include_terminated=False) spark_client.start_stream_to_online_ingestion.assert_called_once_with( feature_table, [], project="default")
def test_new_job_creation(spark_client, feature_table): """ No job existed prior to call """ spark_client.feature_store.list_feature_tables.return_value = [ feature_table ] spark_client.list_jobs.return_value = [] ensure_stream_ingestion_jobs(spark_client, all_projects=True) assert spark_client.start_stream_to_online_ingestion.call_count == 2
def test_new_job_creation(spark_client, feature_table): """ No job existed prior to call """ spark_client.feature_store.list_feature_tables.return_value = [ feature_table ] spark_client.list_jobs.return_value = [] ensure_stream_ingestion_jobs(spark_client, all_projects=True) spark_client.start_stream_to_online_ingestion.assert_called_once_with( feature_table, [], project="default")
def test_restarting_completed_job(spark_client, feature_table): """ Job has succesfully finished on previous try """ job = SimpleStreamingIngestionJob("", "default", feature_table, SparkJobStatus.COMPLETED) spark_client.feature_store.list_feature_tables.return_value = [ feature_table ] spark_client.list_jobs.return_value = [job] ensure_stream_ingestion_jobs(spark_client, all_projects=True) assert spark_client.start_stream_to_online_ingestion.call_count == 2
def test_stopping_running_job(spark_client, feature_table): """ Streaming source was deleted """ new_ft = copy.deepcopy(feature_table) new_ft.stream_source = None job = SimpleStreamingIngestionJob("", feature_table, SparkJobStatus.IN_PROGRESS) spark_client.feature_store.list_feature_tables.return_value = [new_ft] spark_client.list_jobs.return_value = [job] ensure_stream_ingestion_jobs(spark_client, all_projects=True) assert job.get_status() == SparkJobStatus.COMPLETED spark_client.start_stream_to_online_ingestion.assert_not_called()
def test_not_retrying_failed_job(spark_client, feature_table): """ Job has failed on previous try """ job = SimpleStreamingIngestionJob("", feature_table, SparkJobStatus.FAILED) spark_client.feature_store.list_feature_tables.return_value = [ feature_table ] spark_client.list_jobs.return_value = [job] ensure_stream_ingestion_jobs(spark_client, all_projects=True) spark_client.list_jobs.assert_called_once_with(include_terminated=True) assert job.get_status() == SparkJobStatus.FAILED spark_client.start_stream_to_online_ingestion.assert_not_called()
def test_no_changes(spark_client, feature_table): """ Feature Table spec is the same """ job = SimpleStreamingIngestionJob("", feature_table, SparkJobStatus.IN_PROGRESS) spark_client.feature_store.list_feature_tables.return_value = [ feature_table ] spark_client.list_jobs.return_value = [job] ensure_stream_ingestion_jobs(spark_client, all_projects=True) assert job.get_status() == SparkJobStatus.IN_PROGRESS spark_client.start_stream_to_online_ingestion.assert_not_called()
def test_not_cancelling_starting_job(spark_client, feature_table): """ Feature Table spec was updated but previous version is still starting """ new_ft = copy.deepcopy(feature_table) new_ft.stream_source._kafka_options.topic = "new_t" job = SimpleStreamingIngestionJob("", "default", feature_table, SparkJobStatus.STARTING) spark_client.feature_store.list_feature_tables.return_value = [new_ft] spark_client.list_jobs.return_value = [job] ensure_stream_ingestion_jobs(spark_client, all_projects=True) assert job.get_status() == SparkJobStatus.STARTING assert spark_client.start_stream_to_online_ingestion.call_count == 2
def test_update_existing_job(spark_client, feature_table): """ Feature Table spec was updated """ new_ft = copy.deepcopy(feature_table) new_ft.stream_source._kafka_options.topic = "new_t" job = SimpleStreamingIngestionJob("", "default", feature_table, SparkJobStatus.IN_PROGRESS) spark_client.feature_store.list_feature_tables.return_value = [new_ft] spark_client.list_jobs.return_value = [job] ensure_stream_ingestion_jobs(spark_client, all_projects=True) assert job.get_status() == SparkJobStatus.COMPLETED assert spark_client.start_stream_to_online_ingestion.call_count == 2