fix: flake8

Author: evakravi

tests/unit/sagemaker/jumpstart/factory/test_estimator.py

@@ -1,5 +1,18 @@
+# Copyright 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.
+from __future__ import absolute_import
 import pytest
-from unittest.mock import patch, Mock
+from unittest.mock import patch
 from sagemaker.jumpstart.factory.estimator import (
     _add_model_uri_to_kwargs,
     get_model_info_default_kwargs,

tests/unit/sagemaker/jumpstart/test_types.py

@@ -361,7 +361,8 @@ def test_use_training_model_artifact_with_package_uris(self, mock_specs):
             None
         )
         mock_specs.training_model_package_artifact_uris = {
-            "ml.p3.2xlarge": "arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/llama2-13b-e155a2e0347b323fb882f1875851c5d3"
+            "ml.p3.2xlarge": "arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/"
+            "llama2-13b-e155a2e0347b323fb882f1875851c5d3"
         }
 
         result = JumpStartModelSpecs.use_training_model_artifact(mock_specs)