Add unit test for s3_input InputMode

Author: imujjwal96

src/sagemaker/estimator.py

@@ -570,7 +570,7 @@ def start_new(cls, estimator, inputs):
         train_args['metric_definitions'] = estimator.metric_definitions
 
         if isinstance(inputs, s3_input):
-            if inputs.config['InputMode'] is not None:
+            if 'InputMode' in inputs.config:
                 train_args['input_mode'] = inputs.config['InputMode']
 
         if estimator.enable_network_isolation():

tests/unit/test_estimator.py

@@ -315,6 +315,17 @@ def test_augmented_manifest(sagemaker_session):
     assert s3_data_source['AttributeNames'] == ['foo', 'bar']
 
 
+def test_s3_input_mode(sagemaker_session):
+    expected_input_mode = 'Pipe'
+    fw = DummyFramework(entry_point=SCRIPT_PATH, role='DummyRole', sagemaker_session=sagemaker_session,
+                        train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
+                        enable_cloudwatch_metrics=True)
+    fw.fit(inputs=s3_input('s3://mybucket/train_manifest', input_mode=expected_input_mode))
+
+    actual_input_mode = sagemaker_session.method_calls[1][2]['input_mode']
+    assert actual_input_mode == expected_input_mode
+
+
 def test_shuffle_config(sagemaker_session):
     fw = DummyFramework(entry_point=SCRIPT_PATH, role='DummyRole', sagemaker_session=sagemaker_session,
                         train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,