Add unit tests

Author: Dan Choi

src/sagemaker/job.py

@@ -51,14 +51,13 @@ def wait(self):
     def _load_config(inputs, estimator):
         input_config = _Job._format_inputs_to_input_config(inputs)
         role = estimator.sagemaker_session.expand_role(estimator.role)
-        output_config = _Job._prepare_output_config(estimator.output_path,
-                                                            estimator.output_kms_key)
+        output_config = _Job._prepare_output_config(estimator.output_path, estimator.output_kms_key)
         resource_config = _Job._prepare_resource_config(estimator.train_instance_count,
-                                                                estimator.train_instance_type,
-                                                                estimator.train_volume_size)
-        stop_condition = _Job._prepare_stopping_condition(estimator.train_max_run)
+                                                        estimator.train_instance_type,
+                                                        estimator.train_volume_size)
+        stopping_condition = _Job._prepare_stopping_condition(estimator.train_max_run)
 
-        return input_config, role, output_config, resource_config, stop_condition
+        return input_config, role, output_config, resource_config, stopping_condition
 
     @staticmethod
     def _format_inputs_to_input_config(inputs):

tests/unit/test_job.py

@@ -0,0 +1,166 @@
+# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+import pytest
+from mock import Mock
+
+from sagemaker.estimator import Estimator
+from sagemaker.job import _Job
+from sagemaker.session import s3_input
+
+BUCKET_NAME = 's3://mybucket/train'
+S3_OUTPUT_PATH = 's3://bucket/prefix'
+INSTANCE_COUNT = 1
+INSTANCE_TYPE = 'c4.4xlarge'
+VOLUME_SIZE = 1
+MAX_RUNTIME = 1
+ROLE = 'DummyRole'
+IMAGE_NAME = 'fakeimage'
+JOB_NAME = 'fakejob'
+
+
+@pytest.fixture()
+def estimator(sagemaker_session):
+    return Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, VOLUME_SIZE, MAX_RUNTIME,
+                     output_path=S3_OUTPUT_PATH, sagemaker_session=sagemaker_session)
+
+
+@pytest.fixture()
+def job(sagemaker_session):
+    return _Job(sagemaker_session, JOB_NAME)
+
+
+@pytest.fixture()
+def sagemaker_session():
+    boto_mock = Mock(name='boto_session')
+    mock_session = Mock(name='sagemaker_session', boto_session=boto_mock)
+    mock_session.expand_role = Mock(name='expand_role', return_value=ROLE)
+
+    return mock_session
+
+
+def test_load_config(job, estimator):
+    inputs = s3_input(BUCKET_NAME)
+
+    input_config, role, output_config, resource_config, stopping_condition = \
+        job._load_config(inputs, estimator)
+
+    assert input_config[0]['DataSource']['S3DataSource']['S3Uri'] == BUCKET_NAME
+    assert role == ROLE
+    assert output_config['S3OutputPath'] == S3_OUTPUT_PATH
+    assert 'KmsKeyId' not in output_config
+    assert resource_config['InstanceCount'] == INSTANCE_COUNT
+    assert resource_config['InstanceType'] == INSTANCE_TYPE
+    assert resource_config['VolumeSizeInGB'] == VOLUME_SIZE
+    assert stopping_condition['MaxRuntimeInSeconds'] == MAX_RUNTIME
+
+
+def test_format_inputs_to_input_config_string(job):
+    inputs = BUCKET_NAME
+
+    channels = job._format_inputs_to_input_config(inputs)
+
+    assert channels[0]['DataSource']['S3DataSource']['S3Uri'] == inputs
+
+
+def test_format_inputs_to_input_config_s3_input(job):
+    inputs = s3_input(BUCKET_NAME)
+
+    channels = job._format_inputs_to_input_config(inputs)
+
+    assert channels[0]['DataSource']['S3DataSource']['S3Uri'] == inputs.config['DataSource'][
+        'S3DataSource']['S3Uri']
+
+
+def test_format_inputs_to_input_config_dict(job):
+    inputs = {'train': BUCKET_NAME}
+
+    channels = job._format_inputs_to_input_config(inputs)
+
+    assert channels[0]['DataSource']['S3DataSource']['S3Uri'] == inputs['train']
+
+
+def test_format_inputs_to_input_config_exception(job):
+    inputs = 1
+
+    with pytest.raises(ValueError):
+        job._format_inputs_to_input_config(inputs)
+
+
+def test_format_s3_uri_input_string(job):
+    inputs = BUCKET_NAME
+
+    s3_uri_input = job._format_s3_uri_input(inputs)
+
+    assert s3_uri_input.config['DataSource']['S3DataSource']['S3Uri'] == inputs
+
+
+def test_format_s3_uri_input_string_exception(job):
+    inputs = 'mybucket/train'
+
+    with pytest.raises(ValueError):
+        job._format_s3_uri_input(inputs)
+
+
+def test_format_s3_uri_input(job):
+    inputs = s3_input(BUCKET_NAME)
+
+    s3_uri_input = job._format_s3_uri_input(inputs)
+
+    assert s3_uri_input.config['DataSource']['S3DataSource']['S3Uri'] == inputs.config[
+        'DataSource']['S3DataSource']['S3Uri']
+
+
+def test_format_s3_uri_input_exception(job):
+    inputs = 1
+
+    with pytest.raises(ValueError):
+        job._format_s3_uri_input(inputs)
+
+
+def test_prepare_output_config(job):
+    kms_key_id = 'kms_key'
+
+    config = job._prepare_output_config(BUCKET_NAME, kms_key_id)
+
+    assert config['S3OutputPath'] == BUCKET_NAME
+    assert config['KmsKeyId'] == kms_key_id
+
+
+def test_prepare_output_config_kms_key_none(job):
+    s3_path = BUCKET_NAME
+    kms_key_id = None
+
+    config = job._prepare_output_config(s3_path, kms_key_id)
+
+    assert config['S3OutputPath'] == s3_path
+    assert 'KmsKeyId' not in config
+
+
+def test_prepare_resource_config(job):
+    resource_config = job._prepare_resource_config(INSTANCE_COUNT, INSTANCE_TYPE, VOLUME_SIZE)
+
+    assert resource_config['InstanceCount'] == INSTANCE_COUNT
+    assert resource_config['InstanceType'] == INSTANCE_TYPE
+    assert resource_config['VolumeSizeInGB'] == VOLUME_SIZE
+
+
+def test_prepare_stopping_condition(job):
+    max_run = 1
+
+    stopping_condition = job._prepare_stopping_condition(max_run)
+
+    assert stopping_condition['MaxRuntimeInSeconds'] == max_run
+
+
+def test_name(job):
+    assert job.name == JOB_NAME