fix unit tests

Author: laurenyu

src/sagemaker/cli/compatibility/v2/ast_transformer.py

@@ -21,6 +21,7 @@
     modifiers.framework_version.FrameworkVersionEnforcer(),
     modifiers.tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader(),
     modifiers.tf_legacy_mode.TensorBoardParameterRemover(),
+    modifiers.deprecated_params.TensorFlowScriptModeParameterRemover(),
 ]
 
 

src/sagemaker/cli/compatibility/v2/modifiers/__init__.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.cli.compatibility.v2.modifiers import (  # noqa: F401 (imported but unused)
+    deprecated_params,
     framework_version,
     tf_legacy_mode,
 )

src/sagemaker/cli/compatibility/v2/modifiers/deprecated_params.py

@@ -18,7 +18,7 @@
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
 
 
-class TensorFlowScriptModeParamRemover(Modifier):
+class TensorFlowScriptModeParameterRemover(Modifier):
     """A class to remove ``script_mode`` from TensorFlow estimators (because it's the only mode)."""
 
     def node_should_be_modified(self, node):

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_tf_legacy_mode.py

@@ -16,7 +16,7 @@
 
 import pasta
 import pytest
-from mock import patch
+from mock import MagicMock, patch
 
 from sagemaker.cli.compatibility.v2.modifiers import tf_legacy_mode
 from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call
@@ -118,7 +118,9 @@ def test_modify_node_set_image_name_from_args(create_image_uri, boto_session):
     assert expected_string == pasta.dump(node)
 
 
-def test_modify_node_set_hyperparameters():
+@patch("boto3.Session", MagicMock())
+@patch("sagemaker.fw_utils.create_image_uri", return_value=IMAGE_URI)
+def test_modify_node_set_hyperparameters(create_image_uri):
     tf_constructor = """TensorFlow(
         checkpoint_path='s3://foo/bar',
         training_steps=100,
@@ -140,7 +142,9 @@ def test_modify_node_set_hyperparameters():
     assert expected_hyperparameters == _hyperparameters_from_node(node)
 
 
-def test_modify_node_preserve_other_hyperparameters():
+@patch("boto3.Session", MagicMock())
+@patch("sagemaker.fw_utils.create_image_uri", return_value=IMAGE_URI)
+def test_modify_node_preserve_other_hyperparameters(create_image_uri):
     tf_constructor = """sagemaker.tensorflow.TensorFlow(
         training_steps=100,
         evaluation_steps=10,
@@ -164,7 +168,9 @@ def test_modify_node_preserve_other_hyperparameters():
     assert expected_hyperparameters == _hyperparameters_from_node(node)
 
 
-def test_modify_node_prefer_param_over_hyperparameter():
+@patch("boto3.Session", MagicMock())
+@patch("sagemaker.fw_utils.create_image_uri", return_value=IMAGE_URI)
+def test_modify_node_prefer_param_over_hyperparameter(create_image_uri):
     tf_constructor = """sagemaker.tensorflow.TensorFlow(
         training_steps=100,
         requirements_file='source/requirements.txt',

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_tf_script_mode.py

@@ -24,15 +24,15 @@ def test_node_should_be_modified_tf_constructor_script_mode():
         "sagemaker.tensorflow.TensorFlow(script_mode=True)",
     )
 
-    modifier = deprecated_params.TensorFlowScriptModeParamRemover()
+    modifier = deprecated_params.TensorFlowScriptModeParameterRemover()
 
     for constructor in tf_script_mode_constructors:
         node = ast_call(constructor)
         assert modifier.node_should_be_modified(node) is True
 
 
 def test_node_should_be_modified_not_tf_script_mode():
-    modifier = deprecated_params.TensorFlowScriptModeParamRemover()
+    modifier = deprecated_params.TensorFlowScriptModeParameterRemover()
 
     for call in ("TensorFlow()", "random()"):
         node = ast_call(call)
@@ -41,7 +41,7 @@ def test_node_should_be_modified_not_tf_script_mode():
 
 def test_modify_node():
     node = ast_call("TensorFlow(script_mode=True)")
-    modifier = deprecated_params.TensorFlowScriptModeParamRemover()
+    modifier = deprecated_params.TensorFlowScriptModeParameterRemover()
     modifier.modify_node(node)
 
     assert "TensorFlow()" == pasta.dump(node)