def test_batch_prediction_job_done_get(self, get_batch_prediction_job_mock):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME
)
assert bp.done() is False
assert get_batch_prediction_job_mock.call_count == 2
def test_init_batch_prediction_job(self, get_batch_prediction_job_mock):
jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME
)
get_batch_prediction_job_mock.assert_called_once_with(
name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=base._DEFAULT_RETRY
)
def test_batch_prediction_iter_dirs_while_running(self):
"""
Raises RuntimeError since outputs cannot be read while BatchPredictionJob is still running
"""
with pytest.raises(RuntimeError):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME)
bp.iter_outputs()
def test_batch_prediction_iter_dirs_invalid_output_info(self):
"""
Raises NotImplementedError since the BatchPredictionJob's output_info
contains no output GCS directory or BQ dataset.
"""
with pytest.raises(NotImplementedError):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME)
bp.iter_outputs()
def test_batch_prediction_iter_dirs_bq_raises_on_empty(
self, bq_list_rows_mock):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME)
with pytest.raises(RuntimeError) as e:
bp.iter_outputs()
assert e.match(
regexp=("A BigQuery table with predictions was not found,"
" this might be due to errors. Visit http"))
def test_batch_prediction_iter_dirs_gcs(self, storage_list_blobs_mock):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME)
blobs = bp.iter_outputs()
storage_list_blobs_mock.assert_called_once_with(
_TEST_GCS_OUTPUT_INFO.gcs_output_directory, prefix=None)
assert blobs == _TEST_GCS_BLOBS
def test_batch_prediction_iter_dirs_bq(self, bq_list_rows_mock):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME)
bp.iter_outputs()
bq_list_rows_mock.assert_called_once_with(
table=f"{_TEST_BQ_DATASET_ID}.predictions",
max_results=_TEST_BQ_MAX_RESULTS)
def test_batch_prediction_job_status(self, get_batch_prediction_job_mock):
bp = jobs.BatchPredictionJob(
batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME)
# get_batch_prediction() is called again here
bp_job_state = bp.state
assert get_batch_prediction_job_mock.call_count == 2
assert bp_job_state == _TEST_JOB_STATE_SUCCESS
get_batch_prediction_job_mock.assert_called_with(
name=_TEST_BATCH_PREDICTION_JOB_NAME)
