Fix tests

Author: mvsusp

src/sagemaker/estimator.py

@@ -862,16 +862,23 @@ def _stage_user_code_in_s3(self):
         Returns: s3 uri
 
         """
-        if self.code_location is None:
+        local_mode = self.output_path.startswith('file://')
+
+        if self.code_location is None and local_mode:
+            code_bucket = self.sagemaker_session.default_bucket()
+            code_s3_prefix = '{}/source'.format(self._current_job_name)
+            kms_key = None
+
+        elif self.code_location is None:
             code_bucket, _ = parse_s3_url(self.output_path)
             code_s3_prefix = '{}/source'.format(self._current_job_name)
+            kms_key = self.output_kms_key
         else:
             code_bucket, key_prefix = parse_s3_url(self.code_location)
             code_s3_prefix = '/'.join(filter(None, [key_prefix, self._current_job_name, 'source']))
 
-        output_bucket, _ = parse_s3_url(self.output_path)
-
-        kms_key = self.output_kms_key if code_bucket == output_bucket else None
+            output_bucket, _ = parse_s3_url(self.output_path)
+            kms_key = self.output_kms_key if code_bucket == output_bucket else None
 
         return tar_and_upload_dir(session=self.sagemaker_session.boto_session,
                                   bucket=code_bucket,