first round of unit tests for node transformer.

Author: lionsardesai

tests/unit/transformer/test_node_transformer.py

@@ -0,0 +1,48 @@
+#
+#
+# type:ignore
+import ast
+
+import pytest
+from mock import MagicMock
+
+from lineapy.transformer.node_transformer import NodeTransformer
+
+
+class TestNodeTransformer:
+    nt: NodeTransformer
+
+    @pytest.fixture(autouse=True)
+    def before_everything(self):
+        nt = NodeTransformer(
+            "", MagicMock(), MagicMock()  # SourceCodeLocation(0, 0, 0, 0)
+        )
+        assert nt is not None
+        self.nt = nt
+
+    def test_lambda_executes_as_expression(self):
+        self.nt._exec_expression = MagicMock()
+        self.nt._exec_statement = MagicMock()
+
+        # this inits an ast.Module containing one expression whose value is a ast.lambda
+        test_node = ast.parse("lambda x: x + 10")
+        lambda_node = test_node.body[0].value
+        self.nt.generic_visit(lambda_node)
+        self.nt._exec_statement.assert_not_called()
+        self.nt._exec_expression.assert_called_once()
+
+    def test_assign_executes(self):
+        test_node = ast.parse("a = 10")
+        self.nt.visit_Assign = MagicMock()
+        self.nt.visit(test_node.body[0])
+        self.nt.visit_Assign.assert_called_once_with(test_node.body[0])
+
+    def test_assign_calls_tracer_assign(self):
+        self.nt.get_source = MagicMock()
+        test_node = ast.parse("a = 10")
+        tracer = self.nt.tracer
+        self.nt.visit(test_node.body[0])
+        # self.nt.visit_Assign.assert_called_once_with(test_node.body[0])
+        tracer.assign.assert_called_once_with(
+            "a", self.nt.tracer.literal.return_value
+        )

tests/unit/test_source_giver.py renamed to tests/unit/transformer/test_source_giver.py

@@ -6,10 +6,11 @@
 
 import asttokens
 import pytest
-from astpretty import pprint
 
 from lineapy.transformer.source_giver import SourceGiver
 
+# from astpretty import pprint
+
 
 @pytest.mark.parametrize(
     "code,lineno",
@@ -29,14 +30,19 @@ def test_source_giver_adds_end_lineno(code, lineno):
     if sys.version_info >= (3, 8):
         pytest.skip("SourceGiver not invoked for Python 3.8+")
     tree = ast.parse(code)
+    # ensure that the end_lineno is not available and fetching it raises exceptions
     with pytest.raises(AttributeError):
         print(tree.body[0].end_lineno)
 
+    # now we invoke the SourceGiver and add end_linenos in 2 steps - first we run the tree thr asttokens
     asttokens.ASTTokens(code, parse=False, tree=tree)
-    pprint(tree)
-    print(tree.body[0].last_token)
+    # pprint(tree)
+    # print(tree.body[0].last_token)
+    # double check that the line numbers cooked up by asttokens are correct
     assert tree.body[0].last_token.end[0] == lineno
 
+    # and in step 2, run the tree thr SourceGiver and copy the asttokens's token values
+    # so that the tree looks like 3.8+ tree with all the end_linenos etc
     SourceGiver().transform(tree)
     assert tree.body[0].end_lineno == lineno
-    print(tree)
+    # print(tree)