fix: linting, mypy, logical issues for jumpstart models

Author: evakravi

tests/unit/sagemaker/image_uris/test_jumpstart.py

@@ -18,7 +18,6 @@
 import pytest
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
-from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 from sagemaker.jumpstart import constants as sagemaker_constants
 
 
@@ -63,7 +62,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
         sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "*"
     )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         image_uris.retrieve(
             framework=None,
             region="us-west-2",
@@ -73,7 +72,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
             instance_type="ml.p2.xlarge",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         image_uris.retrieve(
             framework=None,
             region="mars-south-1",
@@ -83,7 +82,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
             instance_type="ml.p2.xlarge",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         image_uris.retrieve(
             framework=None,
             region="us-west-2",
@@ -92,7 +91,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
             instance_type="ml.p2.xlarge",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         image_uris.retrieve(
             framework=None,
             region="us-west-2",
@@ -101,7 +100,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
             instance_type="ml.p2.xlarge",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         image_uris.retrieve(
             region="us-west-2",
             framework=None,

tests/unit/sagemaker/test_model_uris.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-from mock.mock import Mock, patch
+from mock.mock import patch
 
 from sagemaker import model_uris
 from sagemaker.jumpstart import constants as sagemaker_constants
@@ -61,35 +61,35 @@ def test_jumpstart_model_uri(patched_get_model_specs):
         sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "*"
     )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             region="us-west-2",
             model_scope="BAD_SCOPE",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             region="mars-south-1",
             model_scope="training",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             model_scope="training",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             model_scope="training",
             model_id="pytorch-ic-mobilenet-v2",

tests/unit/sagemaker/test_script_uris.py

@@ -33,8 +33,8 @@ def test_jumpstart_script_uri(patched_get_model_specs):
         model_version="*",
     )
     assert (
-        uri
-        == f"s3://{get_jumpstart_content_bucket('us-west-2')}/source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz"
+        uri == f"s3://{get_jumpstart_content_bucket('us-west-2')}/"
+        "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz"
     )
     patched_get_model_specs.assert_called_once_with("us-west-2", "pytorch-ic-mobilenet-v2", "*")
 
@@ -47,8 +47,8 @@ def test_jumpstart_script_uri(patched_get_model_specs):
         model_version="*",
     )
     assert (
-        uri
-        == f"s3://{get_jumpstart_content_bucket('us-west-2')}/source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz"
+        uri == f"s3://{get_jumpstart_content_bucket('us-west-2')}/"
+        "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz"
     )
     patched_get_model_specs.assert_called_once_with("us-west-2", "pytorch-ic-mobilenet-v2", "*")
     patched_get_model_specs.reset_mock()
@@ -62,35 +62,35 @@ def test_jumpstart_script_uri(patched_get_model_specs):
         sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "*"
     )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         script_uris.retrieve(
             region="us-west-2",
             script_scope="BAD_SCOPE",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         script_uris.retrieve(
             region="mars-south-1",
             script_scope="training",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         script_uris.retrieve(
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         script_uris.retrieve(
             script_scope="training",
             model_version="*",
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         script_uris.retrieve(
             script_scope="training",
             model_id="pytorch-ic-mobilenet-v2",