fix unit tests

Author: ahsan-z-khan

tests/conftest.py

@@ -14,7 +14,6 @@
 
 import json
 import os
-import pdb
 
 import boto3
 import pytest
@@ -193,8 +192,9 @@ def pytorch_inference_py_version(pytorch_inference_version, request):
 
 @pytest.fixture(scope="module")
 def huggingface_pytorch_training_version(huggingface_training_version):
-    return _huggingface_base_fm_version(huggingface_training_version, "pytorch", "huggingface_training")[0]
-
+    return _huggingface_base_fm_version(
+        huggingface_training_version, "pytorch", "huggingface_training"
+    )[0]
 
 
 @pytest.fixture(scope="module")
@@ -386,14 +386,16 @@ def _huggingface_base_fm_version(huggingface_version, base_fw, fixture_prefix):
 
     for key in list(version_config.keys()):
         if key.startswith(base_fw):
-            base_fw_version = key[len(base_fw):]
+            base_fw_version = key[len(base_fw) :]
             if len(original_version.split(".")) == 2:
                 base_fw_version = ".".join(base_fw_version.split(".")[:-1])
             versions.append(base_fw_version)
     return versions
 
 
-def _generate_huggingface_base_fw_latest_versions(metafunc, fixture_prefix, huggingface_version, base_fw):
+def _generate_huggingface_base_fw_latest_versions(
+    metafunc, fixture_prefix, huggingface_version, base_fw
+):
     versions = _huggingface_base_fm_version(huggingface_version, base_fw, fixture_prefix)
     fixture_name = f"{fixture_prefix}_{base_fw}_latest_version"
 
@@ -414,8 +416,12 @@ def _parametrize_framework_version_fixtures(metafunc, fixture_prefix, config):
         metafunc.parametrize(fixture_name, (latest_version,), scope="session")
 
     if "huggingface" in fixture_prefix:
-        _generate_huggingface_base_fw_latest_versions(metafunc, fixture_prefix, latest_version, "pytorch")
-        _generate_huggingface_base_fw_latest_versions(metafunc, fixture_prefix, latest_version, "tensorflow")
+        _generate_huggingface_base_fw_latest_versions(
+            metafunc, fixture_prefix, latest_version, "pytorch"
+        )
+        _generate_huggingface_base_fw_latest_versions(
+            metafunc, fixture_prefix, latest_version, "tensorflow"
+        )
 
     fixture_name = "{}_latest_py_version".format(fixture_prefix)
     if fixture_name in metafunc.fixturenames:

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -15,7 +15,6 @@
 
 import json
 import os
-import pdb
 
 import pytest
 from mock import MagicMock, Mock, patch
@@ -267,11 +266,12 @@ def test_huggingface(
     assert actual_train_args == expected_train_args
 
 
-def test_attach(sagemaker_session, huggingface_training_version, huggingface_training_pytorch_latest_version):
-    print(huggingface_training_version)
+def test_attach(
+    sagemaker_session, huggingface_training_version, huggingface_pytorch_training_version
+):
     training_image = (
         f"1.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:"
-        f"{huggingface_training_pytorch_latest_version}-"
+        f"{huggingface_pytorch_training_version}-"
         f"transformers{huggingface_training_version}-gpu-py36-cu110-ubuntu18.04"
     )
     returned_job_description = {
@@ -306,7 +306,7 @@ def test_attach(sagemaker_session, huggingface_training_version, huggingface_tra
     assert estimator.latest_training_job.job_name == "neo"
     assert estimator.py_version == "py36"
     assert estimator.framework_version == huggingface_training_version
-    assert estimator.pytorch_version == huggingface_training_pytorch_latest_version
+    assert estimator.pytorch_version == huggingface_pytorch_training_version
     assert estimator.role == "arn:aws:iam::366:role/SageMakerRole"
     assert estimator.instance_count == 1
     assert estimator.max_run == 24 * 60 * 60