def test_main(self): # Mock out all of utils except parser hpo._utils = MagicMock() hpo._utils.add_default_client_arguments = _utils.add_default_client_arguments # Set some static returns hpo._utils.create_hyperparameter_tuning_job.return_value = 'job-name' hpo._utils.get_best_training_job_and_hyperparameters.return_value = 'best_job', { "key_1": "best_hp_1" } hpo._utils.get_image_from_job.return_value = 'training-image' hpo._utils.get_model_artifacts_from_job.return_value = 'model-artifacts' hpo.main(required_args) # Check if correct requests were created and triggered hpo._utils.create_hyperparameter_tuning_job.assert_called() hpo._utils.wait_for_hyperparameter_training_job.assert_called() # Check the file outputs hpo._utils.write_output.assert_has_calls([ call('/tmp/hpo_job_name_output_path', 'job-name'), call('/tmp/model_artifact_url_output_path', 'model-artifacts'), call('/tmp/best_job_name_output_path', 'best_job'), call('/tmp/best_hyperparameters_output_path', {"key_1": "best_hp_1"}, json_encode=True), call('/tmp/training_image_output_path', 'training-image') ])
def test_main_assumes_role(self): # Mock out all of utils except parser hpo._utils = MagicMock() hpo._utils.add_default_client_arguments = _utils.add_default_client_arguments # Set some static returns hpo._utils.create_hyperparameter_tuning_job.return_value = 'job-name' hpo._utils.get_best_training_job_and_hyperparameters.return_value = 'best_job', {"key_1": "best_hp_1"} assume_role_args = required_args + ['--assume_role', 'my-role'] hpo.main(assume_role_args) hpo._utils.get_sagemaker_client.assert_called_once_with('us-west-2', None, assume_role_arn='my-role')
def test_main(self): # Mock out all of utils except parser hpo._utils = MagicMock() hpo._utils.add_default_client_arguments = _utils.add_default_client_arguments # Set some static returns hpo._utils.create_hyperparameter_tuning_job.return_value = 'job-name' hpo._utils.get_best_training_job_and_hyperparameters.return_value = 'best_job', 'best_hyperparameters' hpo._utils.get_image_from_job.return_value = 'training-image' hpo._utils.get_model_artifacts_from_job.return_value = 'model-artifacts' with patch('builtins.open', mock_open()) as file_open: hpo.main(required_args) # Check if correct requests were created and triggered hpo._utils.create_hyperparameter_tuning_job.assert_called() hpo._utils.wait_for_hyperparameter_training_job.assert_called() # Check the file outputs file_open.assert_has_calls([ call('/tmp/hpo_job_name.txt', 'w'), call('/tmp/best_job_name.txt', 'w'), call('/tmp/best_hyperparameters.txt', 'w'), call('/tmp/model_artifact_url.txt', 'w'), call('/tmp/training_image.txt', 'w') ], any_order=True) file_open().write.assert_has_calls([ call('job-name'), call('best_job'), call('"best_hyperparameters"'), call('model-artifacts'), call('training-image'), ], any_order=False)