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