def test_launcher_on_create_endpoint_type_calls_create_endpoint_remote_runner(
self, create_endpoint_remote_runner):
launcher.main(self._input_args)
create_endpoint_remote_runner.assert_called_once_with(
type='CreateEndpoint',
project='test_project',
location='us_central1',
payload='test_payload',
gcp_resources=self._gcp_resources,
executor_input='executor_input')
def test_launcher_on_upload_model_type_calls_upload_model_remote_runner(
self, mock_upload_model_remote_runner):
launcher.main(self._input_args)
mock_upload_model_remote_runner.assert_called_once_with(
type='UploadModel',
project='test_project',
location='us_central1',
payload='test_payload',
gcp_resources=self._gcp_resources,
executor_input='executor_input')
def test_launcher_on_bigquery_query_job_type(
self, bigquery_query_job_remote_runner):
launcher.main(self._input_args)
bigquery_query_job_remote_runner.assert_called_once_with(
type='BigqueryQueryJob',
project='test_project',
location='us_central1',
payload='test_payload',
gcp_resources=self._gcp_resources,
executor_input='executor_input')
def test_launcher_on_custom_job_type_calls_custom_job_remote_runner(
self, mock_custom_job_remote_runner):
job_type = 'CustomJob'
payload = ('{"display_name": "ContainerComponent", "job_spec": '
'{"worker_pool_specs": [{"machine_spec": {"machine_type": '
'"n1-standard-4"}, "replica_count": 1, "container_spec": '
'{"image_uri": "google/cloud-sdk:latest", "command": ["sh", '
'"-c", "set -e -x\\necho \\"$0, this is an output '
'parameter\\"\\n", "{{$.inputs.parameters[\'input_text\']}}", '
'"{{$.outputs.parameters[\'output_value\'].output_file}}"]}}]}}')
input_args = [
'--type', job_type, '--project', self._project, '--location',
self._location, '--payload', payload, '--gcp_resources',
self._gcp_resources, '--extra_arg', 'extra_arg_value'
]
launcher.main(input_args)
mock_custom_job_remote_runner.assert_called_once_with(
type=job_type,
project=self._project,
location=self._location,
payload=payload,
gcp_resources=self._gcp_resources)
def test_launcher_on_batch_prediction_job_type_calls_batch_prediction_job_remote_runner(
self, mock_batch_prediction_job_remote_runner):
job_type = 'BatchPredictionJob'
payload = ('{"batchPredictionJob": {"displayName": '
'"BatchPredictionComponentName", "model": '
'"projects/test/locations/test/models/test-model","inputConfig":'
' {"instancesFormat": "CSV","gcsSource": {"uris": '
'["test_gcs_source"]}}, "outputConfig": {"predictionsFormat": '
'"CSV", "gcsDestination": {"outputUriPrefix": '
'"test_gcs_destination"}}}}')
input_args = [
'--type', job_type, '--project', self._project, '--location',
self._location, '--payload', payload, '--gcp_resources',
self._gcp_resources, '--executor_input', 'executor_input'
]
launcher.main(input_args)
mock_batch_prediction_job_remote_runner.assert_called_once_with(
type=job_type,
project=self._project,
location=self._location,
payload=payload,
gcp_resources=self._gcp_resources,
executor_input='executor_input')