Cleaned up failing tests.

Author: owahab

tests/unit/v2/test_transformer.py

@@ -1,3 +1,5 @@
+from __future__ import absolute_import
+
 import ast
 from sagemaker.tools.compatibility.v2.ast_transformer import ASTTransformer
 import pasta
@@ -10,20 +12,15 @@ def test_code_needs_transform():
 m = MXNet()
 sagemaker.mxnet.MXNet()
 """
+
     transformer_class = ASTTransformer()
-    rewrite = transformer_class.visit(
-        ast.parse(
-            simple
-        )
-    )
+    rewrite = transformer_class.visit(ast.parse(simple))
     expected = """TensorFlow(entry_point='foo.py', framework_version='1.11.0')
 sagemaker.tensorflow.TensorFlow(framework_version='1.11.0')
 m = MXNet(framework_version='1.2.0')
 sagemaker.mxnet.MXNet(framework_version='1.2.0')\n"""
 
-    assert pasta.dump(
-        rewrite
-    ) == expected
+    assert pasta.dump(rewrite) == expected
 
 
 def test_code_does_not_need_transform():
@@ -32,16 +29,10 @@ def test_code_does_not_need_transform():
 m = MXNet(framework_version='1.2.0')
 sagemaker.mxnet.MXNet(framework_version='1.2.0')\n"""
     transformer_class = ASTTransformer()
-    rewrite = transformer_class.visit(
-        ast.parse(
-            simple
-        )
-    )
+    rewrite = transformer_class.visit(ast.parse(simple))
     expected = """TensorFlow(entry_point='foo.py', framework_version='1.11.0')
 sagemaker.tensorflow.TensorFlow(framework_version='1.11.0')
 m = MXNet(framework_version='1.2.0')
 sagemaker.mxnet.MXNet(framework_version='1.2.0')\n"""
 
-    assert pasta.dump(
-        rewrite
-    ) == expected
+    assert pasta.dump(rewrite) == expected