Revert "first pass"

This reverts commit b507770.

Author: KRRT7

codeflash/cli_cmds/cmd_init.py

@@ -472,6 +472,8 @@ def collect_setup_info() -> CLISetupInfo:
     for choice in test_framework_choices:
         if choice == "pytest":
             framework_choices.append(("🧪 pytest", "pytest"))
+        elif choice == "unittest":
+            framework_choices.append(("🐍 unittest", "unittest"))
 
     framework_questions = [
         inquirer.List(
@@ -615,6 +617,31 @@ def detect_test_framework_from_config_files(curdir: Path) -> Optional[str]:
     return test_framework
 
 
+def detect_test_framework_from_test_files(tests_root: Path) -> Optional[str]:
+    test_framework = None
+    # Check if any python files contain a class that inherits from unittest.TestCase
+    for filename in tests_root.iterdir():
+        if filename.suffix == ".py":
+            with filename.open(encoding="utf8") as file:
+                contents = file.read()
+                try:
+                    node = ast.parse(contents)
+                except SyntaxError:
+                    continue
+                if any(
+                    isinstance(item, ast.ClassDef)
+                    and any(
+                        (isinstance(base, ast.Attribute) and base.attr == "TestCase")
+                        or (isinstance(base, ast.Name) and base.id == "TestCase")
+                        for base in item.bases
+                    )
+                    for item in node.body
+                ):
+                    test_framework = "unittest"
+                    break
+    return test_framework
+
+
 def check_for_toml_or_setup_file() -> str | None:
     click.echo()
     click.echo("Checking for pyproject.toml or setup.py…\r", nl=False)
@@ -1298,7 +1325,26 @@ def sorter(arr: Union[List[int],List[float]]) -> Union[List[int],List[float]]:
                 arr[j + 1] = temp
     return arr
 """
-    bubble_sort_test_content = f"""from {Path(args.module_root).name}.bubble_sort import sorter
+    if args.test_framework == "unittest":
+        bubble_sort_test_content = f"""import unittest
+from {os.path.basename(args.module_root)}.bubble_sort import sorter # Keep usage of os.path.basename to avoid pathlib potential incompatibility https://github.com/codeflash-ai/codeflash/pull/1066#discussion_r1801628022
+
+class TestBubbleSort(unittest.TestCase):
+    def test_sort(self):
+        input = [5, 4, 3, 2, 1, 0]
+        output = sorter(input)
+        self.assertEqual(output, [0, 1, 2, 3, 4, 5])
+
+        input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
+        output = sorter(input)
+        self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
+
+        input = list(reversed(range(100)))
+        output = sorter(input)
+        self.assertEqual(output, list(range(100)))
+"""  # noqa: PTH119
+    elif args.test_framework == "pytest":
+        bubble_sort_test_content = f"""from {Path(args.module_root).name}.bubble_sort import sorter
 
 def test_sort():
     input = [5, 4, 3, 2, 1, 0]

codeflash/code_utils/concolic_utils.py

@@ -28,11 +28,44 @@ def _transform_assert_line(self, line: str) -> Optional[str]:
             expression = expression.rstrip(",;")
             return f"{indent}{expression}"
 
+        unittest_match = self.unittest_re.match(line)
+        if unittest_match:
+            indent, assert_method, args = unittest_match.groups()
+
+            if args:
+                arg_parts = self._split_top_level_args(args)
+                if arg_parts and arg_parts[0]:
+                    return f"{indent}{arg_parts[0]}"
+
         return None
 
+    def _split_top_level_args(self, args_str: str) -> list[str]:
+        result = []
+        current = []
+        depth = 0
+
+        for char in args_str:
+            if char in "([{":
+                depth += 1
+                current.append(char)
+            elif char in ")]}":
+                depth -= 1
+                current.append(char)
+            elif char == "," and depth == 0:
+                result.append("".join(current).strip())
+                current = []
+            else:
+                current.append(char)
+
+        if current:
+            result.append("".join(current).strip())
+
+        return result
+
     def __init__(self) -> None:
         # Pre-compiling regular expressions for faster execution
         self.assert_re = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$")
+        self.unittest_re = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$")
 
 
 def clean_concolic_tests(test_suite_code: str) -> str:

codeflash/code_utils/instrument_existing_tests.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import platform
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -318,6 +319,7 @@ def iter_ast_calls(node):  # noqa: ANN202, ANN001
         return return_statement
 
     def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
+        # TODO: Ensure that this class inherits from unittest.TestCase. Don't modify non unittest.TestCase classes.
         for inner_node in ast.walk(node):
             if isinstance(inner_node, ast.FunctionDef):
                 self.visit_FunctionDef(inner_node, node.name)
@@ -327,6 +329,17 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
     def visit_FunctionDef(self, node: ast.FunctionDef, test_class_name: str | None = None) -> ast.FunctionDef:
         if node.name.startswith("test_"):
             did_update = False
+            if self.test_framework == "unittest" and platform.system() != "Windows":
+                # Only add timeout decorator on non-Windows platforms
+                # Windows doesn't support SIGALRM signal required by timeout_decorator
+
+                node.decorator_list.append(
+                    ast.Call(
+                        func=ast.Name(id="timeout_decorator.timeout", ctx=ast.Load()),
+                        args=[ast.Constant(value=15)],
+                        keywords=[],
+                    )
+                )
             i = len(node.body) - 1
             while i >= 0:
                 line_node = node.body[i]
@@ -492,6 +505,25 @@ def __init__(
             self.class_name = function.top_level_parent_name
 
     def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
+        # Add timeout decorator for unittest test classes if needed
+        if self.test_framework == "unittest":
+            timeout_decorator = ast.Call(
+                func=ast.Name(id="timeout_decorator.timeout", ctx=ast.Load()),
+                args=[ast.Constant(value=15)],
+                keywords=[],
+            )
+            for item in node.body:
+                if (
+                    isinstance(item, ast.FunctionDef)
+                    and item.name.startswith("test_")
+                    and not any(
+                        isinstance(d, ast.Call)
+                        and isinstance(d.func, ast.Name)
+                        and d.func.id == "timeout_decorator.timeout"
+                        for d in item.decorator_list
+                    )
+                ):
+                    item.decorator_list.append(timeout_decorator)
         return self.generic_visit(node)
 
     def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef:
@@ -510,6 +542,25 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
     def _process_test_function(
         self, node: ast.AsyncFunctionDef | ast.FunctionDef
     ) -> ast.AsyncFunctionDef | ast.FunctionDef:
+        # Optimize the search for decorator presence
+        if self.test_framework == "unittest":
+            found_timeout = False
+            for d in node.decorator_list:
+                # Avoid isinstance(d.func, ast.Name) if d is not ast.Call
+                if isinstance(d, ast.Call):
+                    f = d.func
+                    # Avoid attribute lookup if f is not ast.Name
+                    if isinstance(f, ast.Name) and f.id == "timeout_decorator.timeout":
+                        found_timeout = True
+                        break
+            if not found_timeout:
+                timeout_decorator = ast.Call(
+                    func=ast.Name(id="timeout_decorator.timeout", ctx=ast.Load()),
+                    args=[ast.Constant(value=15)],
+                    keywords=[],
+                )
+                node.decorator_list.append(timeout_decorator)
+
         # Initialize counter for this test function
         if node.name not in self.async_call_counter:
             self.async_call_counter[node.name] = 0
@@ -664,6 +715,8 @@ def inject_async_profiling_into_existing_test(
 
     # Add necessary imports
     new_imports = [ast.Import(names=[ast.alias(name="os")])]
+    if test_framework == "unittest":
+        new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
 
     tree.body = [*new_imports, *tree.body]
     return True, sort_imports(ast.unparse(tree), float_to_top=True)
@@ -709,6 +762,8 @@ def inject_profiling_into_existing_test(
                 ast.Import(names=[ast.alias(name="dill", asname="pickle")]),
             ]
         )
+    if test_framework == "unittest" and platform.system() != "Windows":
+        new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
     additional_functions = [create_wrapper_function(mode)]
 
     tree.body = [*new_imports, *additional_functions, *tree.body]