fix: check ast node on later renamers for cli v2 updater (#1848)

Author: metrizable

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

@@ -112,8 +112,10 @@ def node_should_be_modified(self, node):
         Returns:
             bool: If the import statement imports ``get_image_uri`` from the correct module.
         """
-        return node.module in GET_IMAGE_URI_NAMESPACES and any(
-            name.name == GET_IMAGE_URI_NAME for name in node.names
+        return (
+            node is not None
+            and node.module in GET_IMAGE_URI_NAMESPACES
+            and any(name.name == GET_IMAGE_URI_NAME for name in node.names)
         )
 
     def modify_node(self, node):

tests/unit/sagemaker/cli/compatibility/v2/test_ast_transformer.py

@@ -0,0 +1,48 @@
+# Copyright 2020 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 textwrap
+
+import pasta
+import pytest
+
+from sagemaker.cli.compatibility.v2.ast_transformer import ASTTransformer
+
+
+@pytest.fixture
+def input_code():
+    return textwrap.dedent(
+        """
+        from sagemaker.predictor import csv_serializer
+
+        csv_serializer.__doc__
+        """
+    )
+
+
+@pytest.fixture
+def output_code():
+    return textwrap.dedent(
+        """
+        from sagemaker import serializers
+
+        serializers.CSVSerializer().__doc__
+        """
+    )
+
+
+def test_simple_script(input_code, output_code):
+    input_ast = pasta.parse(input_code)
+    output_ast = ASTTransformer().visit(input_ast)
+    assert pasta.dump(output_ast) == output_code