def test_job_id_parameters_override(self, mock_load_json,
mock_get_current_time):
mock_load_json.return_value = _load_test_data('pipeline_job.json')
mock_get_current_time.return_value = datetime.date(2020, 10, 28)
api_client = client.AIPlatformClient(
project_id='test-project', region='us-central1')
with mock.patch.object(
api_client, '_submit_job', autospec=True) as mock_submit:
api_client.create_run_from_job_spec(
job_spec_path='path/to/pipeline_job.json',
job_id='my-new-id',
pipeline_root='gs://bucket/new-blob',
parameter_values={
'text': 'Hello test!',
'list': [1, 2, 3],
})
golden = _load_test_data('pipeline_job.json')
golden['name'] = ('projects/test-project/locations/us-central1/'
'pipelineJobs/my-new-id')
golden['displayName'] = 'my-new-id'
golden['runtimeConfig'][
'gcsOutputDirectory'] = 'gs://bucket/new-blob'
golden['runtimeConfig']['parameters']['text'] = {
'stringValue': 'Hello test!'
}
golden['runtimeConfig']['parameters']['list'] = {
'stringValue': '[1, 2, 3]'
}
mock_submit.assert_called_once_with(
job_spec=golden, job_id='my-new-id')
def test_client_init_with_defaults(self):
api_client = client.AIPlatformClient(
project_id='test-project', region='us-central1')
self.assertEqual(api_client._project_id, 'test-project')
self.assertEqual(api_client._region, 'us-central1')
self.assertEqual(api_client._endpoint,
'us-central1-aiplatform.googleapis.com')
def test_get_job_success(self, mock_build_client, mock_get_gcp_credential,
mock_get_current_time):
mock_get_current_time.return_value = datetime.date(2020, 10, 28)
mock_build_client.return_value = _MockClient()
api_client = client.AIPlatformClient(
project_id='test-project', region='us-central1')
self.assertEqual(_EXPECTED_GET_RESPONSE,
api_client.get_job(job_id='test-job'))
mock_get_gcp_credential.assert_called_once()
def test_create_run_from_pipeline_job(self, mock_load_json,
mock_get_current_time):
mock_load_json.return_value = _load_test_data('pipeline_job.json')
mock_get_current_time.return_value = datetime.date(2020, 10, 28)
api_client = client.AIPlatformClient(
project_id='test-project', region='us-central1')
with mock.patch.object(
api_client, '_submit_job', autospec=True) as mock_submit:
api_client.create_run_from_job_spec(
job_spec_path='path/to/pipeline_job.json')
golden = _load_test_data('pipeline_job.json')
mock_submit.assert_called_once_with(
job_spec=golden, job_id='sample-test-pipeline-20201028000000')
def test_disable_caching(self, mock_load_json, mock_get_current_time):
mock_load_json.return_value = _load_test_data('pipeline_job.json')
mock_get_current_time.return_value = datetime.date(2020, 10, 28)
api_client = client.AIPlatformClient(
project_id='test-project', region='us-central1')
with mock.patch.object(
api_client, '_submit_job', autospec=True) as mock_submit:
api_client.create_run_from_job_spec(
job_spec_path='path/to/pipeline_job.json', enable_caching=False)
golden = _load_test_data('pipeline_job.json')
golden = json.loads(
json.dumps(golden).replace('"enableCache": true',
'"enableCache": false'))
mock_submit.assert_called_once_with(
job_spec=golden, job_id='sample-test-pipeline-20201028000000')
def test_advanced_settings(self, mock_load_json, mock_get_current_time):
mock_load_json.return_value = _load_test_data('pipeline_job.json')
mock_get_current_time.return_value = datetime.date(2020, 10, 28)
api_client = client.AIPlatformClient(
project_id='test-project', region='us-central1')
with mock.patch.object(
api_client, '_submit_job', autospec=True) as mock_submit:
api_client.create_run_from_job_spec(
job_spec_path='path/to/pipeline_job.json',
cmek='custom-key',
service_account='custom-sa',
network='custom-network')
golden = _load_test_data('pipeline_job.json')
golden['encryptionSpec'] = {'kmsKeyName': 'custom-key'}
golden['serviceAccount'] = 'custom-sa'
golden['network'] = 'custom-network'
mock_submit.assert_called_once_with(
job_spec=golden, job_id='sample-test-pipeline-20201028000000')