def test_dictionary_generated(table, staging_bucket, result_column, save_mode,
options, expected_dict, expected_valid):
bq_sink = BigQuerySink(table, staging_bucket, result_column, save_mode,
options)
if expected_valid:
assert bq_sink.to_dict() == expected_dict
else:
with pytest.raises(ValueError):
bq_sink.to_dict()
def test_create_prediction_job(self, version):
job_1.status = "completed"
responses.add("POST", '/v1/models/1/versions/1/jobs',
body=json.dumps(job_1.to_dict()),
status=200,
content_type='application/json')
bq_src = BigQuerySource(table="project.dataset.source_table",
features=["feature_1", "feature2"],
options={"key": "val"})
bq_sink = BigQuerySink(table="project.dataset.result_table",
result_column="prediction",
save_mode=SaveMode.OVERWRITE,
staging_bucket="gs://test",
options={"key": "val"})
job_config = PredictionJobConfig(source=bq_src,
sink=bq_sink,
service_account_name="my-service-account",
result_type=ResultType.INTEGER)
j = version.create_prediction_job(job_config=job_config)
assert j.status == JobStatus.COMPLETED
assert j.id == job_1.id
assert j.error == job_1.error
assert j.name == job_1.name
actual_req = json.loads(responses.calls[0].request.body)
assert actual_req["config"]["job_config"]["bigquery_source"] == bq_src.to_dict()
assert actual_req["config"]["job_config"]["bigquery_sink"] == bq_sink.to_dict()
assert actual_req["config"]["job_config"]["model"]["result"]["type"] == ResultType.INTEGER.value
assert actual_req["config"]["job_config"]["model"]["uri"] == f"{version.artifact_uri}/model"
assert actual_req["config"]["job_config"]["model"]["type"] == ModelType.PYFUNC_V2.value.upper()
assert actual_req["config"]["service_account_name"] == "my-service-account"
def test_create_prediction_job_with_retry_success(self, version):
job_1.status = "pending"
responses.add("POST",
'/v1/models/1/versions/1/jobs',
body=json.dumps(job_1.to_dict()),
status=200,
content_type='application/json')
# Patch the method as currently it is not supported in the library
# https://github.com/getsentry/responses/issues/135
def _find_match(self, request):
for match in self._urls:
if request.method == match['method'] and \
self._has_url_match(match, request.url):
return match
def _find_match_patched(self, request):
for index, match in enumerate(self._urls):
if request.method == match['method'] and \
self._has_url_match(match, request.url):
if request.method == "GET" and request.url == "/v1/models/1/versions/1/jobs/1":
return self._urls.pop(index)
else:
return match
responses._find_match = types.MethodType(_find_match_patched,
responses)
for i in range(4):
responses.add("GET",
'/v1/models/1/versions/1/jobs/1',
body=json.dumps(job_1.to_dict()),
status=500,
content_type='application/json')
job_1.status = "completed"
responses.add("GET",
'/v1/models/1/versions/1/jobs/1',
body=json.dumps(job_1.to_dict()),
status=200,
content_type='application/json')
bq_src = BigQuerySource(table="project.dataset.source_table",
features=["feature_1", "feature2"],
options={"key": "val"})
bq_sink = BigQuerySink(table="project.dataset.result_table",
result_column="prediction",
save_mode=SaveMode.OVERWRITE,
staging_bucket="gs://test",
options={"key": "val"})
job_config = PredictionJobConfig(
source=bq_src,
sink=bq_sink,
service_account_name="my-service-account",
result_type=ResultType.INTEGER)
j = version.create_prediction_job(job_config=job_config)
assert j.status == JobStatus.COMPLETED
assert j.id == job_1.id
assert j.error == job_1.error
assert j.name == job_1.name
actual_req = json.loads(responses.calls[0].request.body)
assert actual_req["config"]["job_config"][
"bigquery_source"] == bq_src.to_dict()
assert actual_req["config"]["job_config"][
"bigquery_sink"] == bq_sink.to_dict()
assert actual_req["config"]["job_config"]["model"]["result"][
"type"] == ResultType.INTEGER.value
assert actual_req["config"]["job_config"]["model"][
"uri"] == f"{version.artifact_uri}/model"
assert actual_req["config"]["job_config"]["model"][
"type"] == ModelType.PYFUNC_V2.value.upper()
assert actual_req["config"][
"service_account_name"] == "my-service-account"
assert len(responses.calls) == 6
# unpatch
responses._find_match = types.MethodType(_find_match, responses)