add returns

Author: Dan Choi

src/sagemaker/local/image.py

@@ -659,11 +659,11 @@ def _write_json_file(filename, content):
 def _ecr_login_if_needed(boto_session, image):
     # Only ECR images need login
     if not ('dkr.ecr' in image and 'amazonaws.com' in image):
-        return
+        return False
 
     # do we have the image?
     if _check_output('docker images -q %s' % image).strip():
-        return
+        return False
 
     if not boto_session:
         raise RuntimeError('A boto session is required to login to ECR.'

tests/unit/test_image.py

@@ -537,20 +537,22 @@ def test_prepare_serving_volumes_with_local_model(get_data_source_instance, sage
 
 def test_ecr_login_non_ecr():
     session_mock = Mock()
-    sagemaker.local.image._ecr_login_if_needed(session_mock, 'ubuntu')
+    result = sagemaker.local.image._ecr_login_if_needed(session_mock, 'ubuntu')
 
     session_mock.assert_not_called()
+    assert result is False
 
 
 @patch('sagemaker.local.image._check_output', return_value='123451324')
 def test_ecr_login_image_exists(_check_output):
     session_mock = Mock()
 
     image = '520713654638.dkr.ecr.us-east-1.amazonaws.com/image-i-have:1.0'
-    sagemaker.local.image._ecr_login_if_needed(session_mock, image)
+    result = sagemaker.local.image._ecr_login_if_needed(session_mock, image)
 
     session_mock.assert_not_called()
     _check_output.assert_called()
+    assert result is False
 
 
 @patch('subprocess.check_output', return_value=''.encode('utf-8'))