Major WIP.

Author: owahab

tests/unit/v2/__init__.py

@@ -0,0 +1,4 @@
+TensorFlow(entry_point="foo.py")
+sagemaker.tensorflow.TensorFlow()
+m = MXNet()
+sagemaker.mxnet.MXNet()

tests/unit/v2/samples/simple.txt

@@ -0,0 +1,13 @@
+import unittest
+
+
+class FrameworkVersion(unittest.TestCase):
+    def setUp(self) -> None:
+        pass
+
+    def test_something(self):
+        self.assertEqual(True, False)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/unit/v2/test_framework_version.py

@@ -0,0 +1,32 @@
+import ast
+import unittest
+
+from tests.unit.v2.utils import get_sample_file
+from tools.compatibility.v2.ast_transformer import ASTTransformer
+import pasta
+
+
+class TransformerTest(unittest.TestCase):
+    def setUp(self) -> None:
+        self.transformer_class = ASTTransformer()
+
+    def test_simple_transform(self):
+        sample = get_sample_file('simple.txt')
+        rewrite = self.transformer_class.visit(
+            ast.parse(
+                sample
+            )
+        )
+
+        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"""
+
+        self.assertEqual(pasta.dump(rewrite), expected)
+
+
+if __name__ == '__main__':
+    unittest.main()
+
+

tests/unit/v2/test_transformer.py

@@ -0,0 +1,9 @@
+from os.path import join
+
+SAMPLES_DIRECTORY = "/Users/owahab/Desktop/personal/sagemaker-python-sdk/tests/unit/v2/samples/"
+
+
+def get_sample_file(filename):
+    file_path = join(SAMPLES_DIRECTORY, filename)
+    with open(file_path) as file_content:
+        return file_content.read()

tests/unit/v2/utils.py

@@ -15,7 +15,7 @@
 
 import ast
 
-from modifiers import framework_version
+from tools.compatibility.v2.modifiers import framework_version
 
 FUNCTION_CALL_MODIFIERS = [framework_version.FrameworkVersionEnforcer()]
 

tools/compatibility/v2/ast_transformer.py

@@ -15,7 +15,7 @@
 
 import ast
 
-from modifiers.modifier import Modifier
+from tools.compatibility.v2.modifiers.modifier import Modifier
 
 FRAMEWORK_DEFAULTS = {
     "Chainer": "4.1.0",