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)