Sweep: Write unit tests for code_to_data_dependency_graph.py and var_ddg.py (✓ Sandbox Passed)

code_to_data_dependency_graph.py

@@ -2,6 +2,17 @@
 
 
 def get_deps(code):
+    """Extracts dependencies from a block of code.
+
+    This function analyzes the given code to extract its Data Dependency Graph (DDG) and the declaration line number map. It parses the code,
+    identifies variables, and tracks the dependencies between them.
+
+    Args:
+        code (str): The source code to analyze.
+
+    Returns:
+        tuple: A tuple containing the declaration line number map (dict) and the DDG (dict).
+    """
     body = ast.parse(code)
     _, statements = next(ast.iter_fields(body))
 
@@ -63,6 +74,17 @@ def visit_FunctionDef(self, node):
         return fn_nodes
 
     def recursive_ddg(self, fn_root_node):
+        """Recursively builds the DDG for a given function root node.
+
+        This method visits each assignment within the function and constructs the DDG by tracking variables and their dependencies. It handles
+        embedded function definitions and variable scopes within the method.
+
+        Args:
+            fn_root_node (ast.FunctionDef): The root AST node of the function to analyze.
+
+        Returns:
+            dict: The constructed DDG.
+        """
         ddg = {}
         self_edge_set = set()
 

tests/test_code_to_data_dependency_graph.py

@@ -0,0 +1,65 @@
+import unittest
+
+from code_to_data_dependency_graph import MethodLevelDDGs, get_deps
+
+
+class TestGetDepsFunction(unittest.TestCase):
+
+    def test_simple_assignment(self):
+        code = "x = 1"
+        decl_map, ddg = get_deps(code)
+        self.assertEqual(decl_map, {'x': 0})
+        self.assertEqual(ddg, {'x': set()})
+
+    def test_multiple_assignments_single_statement(self):
+        code = "x, y = 1, 2"
+        decl_map, ddg = get_deps(code)
+        self.assertEqual(decl_map, {'x': 0, 'y': 0})
+        self.assertEqual(ddg, {'x': set(), 'y': set()})
+
+    def test_assignment_with_function_call(self):
+        code = "def f(a): return a\nx = f(1)"
+        decl_map, ddg = get_deps(code)
+        self.assertEqual(decl_map, {'f': 0, 'x': 1})
+        self.assertTrue('f' in ddg['x'])
+
+class TestMethodLevelDDGs(unittest.TestCase):
+
+    def test_simple_method_dependency(self):
+        code = """
+def a(x):
+    return x
+
+def b(y):
+    return a(y)
+"""
+        ml_ddgs = MethodLevelDDGs(code)
+        ddgs = ml_ddgs.fn_ddgs(code)
+        self.assertTrue('a' in ddgs['b'])
+
+    def test_nested_function_definitions(self):
+        code = """
+def outer(x):
+    def inner(y):
+        return y
+    return inner(x)
+"""
+        ml_ddgs = MethodLevelDDGs(code)
+        ddgs = ml_ddgs.fn_ddgs(code)
+        self.assertTrue('inner' in ddgs['_global_'])
+        self.assertTrue('x' in ddgs['outer'])
+
+    def test_inter_method_dependencies(self):
+        code = """
+def first():
+    return 42
+
+def second():
+    return first()
+"""
+        ml_ddgs = MethodLevelDDGs(code)
+        ddgs = ml_ddgs.fn_ddgs(code)
+        self.assertTrue('first' in ddgs['second'])
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_var_ddg.py

@@ -0,0 +1,52 @@
+import unittest
+from io import StringIO
+from unittest.mock import patch
+
+from var_ddg import MyVisitor
+
+
+class TestMyVisitor(unittest.TestCase):
+
+    @patch("builtins.open", new_callable=lambda: lambda *args, **kwargs: StringIO("x = 1\ny = x + 2"))
+    def test_simple_dependency(self, mock_open):
+        visitor = MyVisitor("dummy_path")
+        visitor.construct_ddg()
+        ddg = visitor.get_function_level_ddg()
+        self.assertIn('_global_', ddg)
+        self.assertIn(('y', 'x'), ddg['_global_'])
+
+    @patch("builtins.open", new_callable=lambda: lambda *args, **kwargs: StringIO("for i in range(5):\n    x = i"))
+    def test_loop_dependency(self, mock_open):
+        visitor = MyVisitor("dummy_path")
+        visitor.construct_ddg()
+        ddg = visitor.get_function_level_ddg()
+        self.assertIn('_global_', ddg)
+        self.assertIn(('x', 'i'), ddg['_global_'])
+
+    @patch("builtins.open", new_callable=lambda: lambda *args, **kwargs: StringIO("if condition:\n    x = 42\nelse:\n    x = 0"))
+    def test_conditional_dependency(self, mock_open):
+        visitor = MyVisitor("dummy_path")
+        visitor.construct_ddg()
+        ddg = visitor.get_function_level_ddg()
+        self.assertIn('_global_', ddg)
+        self.assertIn(('x', 'condition'), ddg['_global_'])
+
+    @patch("builtins.open", new_callable=lambda: lambda *args, **kwargs: StringIO("def foo():\n    x = 1\n    return x\ny = foo()"))
+    def test_function_call_dependency(self, mock_open):
+        visitor = MyVisitor("dummy_path")
+        visitor.construct_ddg()
+        ddg = visitor.get_function_level_ddg()
+        self.assertIn('_global_', ddg)
+        self.assertIn(('y', 'foo'), ddg['_global_'])
+
+    @patch("builtins.open", new_callable=lambda: lambda *args, **kwargs: StringIO("x = 1\ny = x + 2\ndef foo():\n    z = y"))
+    def test_construct_ddg(self, mock_open):
+        visitor = MyVisitor("dummy_path")
+        visitor.construct_ddg()
+        ddg = visitor.get_function_level_ddg()
+        self.assertIn('_global_', ddg)
+        self.assertIn(('y', 'x'), ddg['_global_'])
+        self.assertIn(('z', 'y'), ddg['foo'])
+
+if __name__ == '__main__':
+    unittest.main()

var_ddg.py

@@ -47,6 +47,13 @@ def ast_visit(self, node, level=0):
 
             # This ast.walk() call in the loop causes the complexity to be O(n^2)
             for descendant in ast.walk(node.value):
+        """
+        Visits nodes in the AST recursively to construct the DDG.
+
+        Args:
+            node (ast.AST): The current node being visited.
+            level (int): The current level of depth in the AST.
+        """
                 if isinstance(descendant, ast.Name):
                     depends_on.append(descendant)
             for var in lhs_ids: