Merge pull request #1399 from codeflash-ai/fix-new-test-instrumentation-kwargs

Fix django test instrumentation

django/aiservice/testgen/instrumentation/instrument_new_tests.py

@@ -229,7 +229,7 @@ class InjectPerfAndLogging(ast.NodeTransformer):
                         ast.Constant(value=self.only_function_name),
                         ast.Constant(value=index),
                         ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
-                        *([test_node.args] if self.mode == TestingMode.PERFORMANCE else []),
+                        *(test_node.args if self.mode == TestingMode.PERFORMANCE else []),
                     ],
                     keywords=[
                         ast.keyword(

django/aiservice/testgen/instrumentation/tests/test_instrument_generated_tests.py

@@ -388,6 +388,228 @@ def test_large_list():
     assert result == expected_perfonly_result
 
 
+def test_instrument_test_source3() -> None:
+    test_source = """from __future__ import annotations
+
+import typing as t
+from unittest.mock import Mock
+
+# imports
+import pytest  # used for our unit tests
+from src.click.core import CommandCollection
+
+
+class Command:
+    pass  # Placeholder for the Command class
+
+class Context:
+    pass  # Placeholder for the Context class
+from src.click.core import CommandCollection
+
+# unit tests
+
+@pytest.fixture
+def mock_context():
+    return Mock(spec=Context)
+
+@pytest.fixture
+def mock_command():
+    return Mock(spec=Command)
+
+@pytest.fixture
+def mock_group():
+    group = Mock(spec=Group)
+    group.get_command.return_value = None
+    return group
+
+def test_command_exists_in_current_collection(mock_context, mock_command):
+    # Setup
+    collection = CommandCollection()
+    collection.get_command = Mock(return_value=mock_command)
+
+    # Test
+    codeflash_output = collection.get_command(mock_context, "start")
+
+def test_command_exists_in_source_group(mock_context, mock_command, mock_group):
+    # Setup
+    mock_group.get_command.return_value = mock_command
+    collection = CommandCollection(sources=[mock_group])
+
+    # Test
+    codeflash_output = collection.get_command(mock_context, option="stop")
+
+def test_command_does_not_exist(mock_context, mock_group):
+    # Setup
+    collection = CommandCollection(sources=[mock_group])
+    args = ["start", "stop", "restart"]
+    kwargs = {"option2": "stop"}
+
+    # Test
+    codeflash_output = collection.get_command(mock_context, *args, option="restart", **kwargs)
+
+def test_empty_command_name(mock_context, mock_group):
+    # Setup
+    collection = CommandCollection(sources=[mock_group])
+
+    # Test
+    codeflash_output = collection.get_command(mock_context)
+
+def test_none_as_command_name(mock_context, mock_group):
+    # Setup
+    collection = CommandCollection(sources=[mock_group])
+
+    # Test
+    codeflash_output = collection.get_command(mock_context, sources=[mock_group])
+"""
+
+    expected = format_and_float_to_top(
+        """from __future__ import annotations
+
+import gc
+import os
+import time
+import typing as t
+from typing import Any, Callable, Optional
+from unittest.mock import Mock
+
+import pytest
+from code_to_optimize.bubble_sort_in_class import CommandCollection
+
+"""
+        + perf_logging_code
+        + """
+
+
+class Command:
+    pass
+
+
+class Context:
+    pass
+
+
+@pytest.fixture
+def mock_context():
+    return Mock(spec=Context)
+
+
+@pytest.fixture
+def mock_command():
+    return Mock(spec=Command)
+
+
+@pytest.fixture
+def mock_group():
+    group = Mock(spec=Group)
+    group.get_command.return_value = None
+    return group
+
+
+def test_command_exists_in_current_collection(mock_context, mock_command):
+    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
+    collection = CommandCollection()
+    collection.get_command = Mock(return_value=mock_command)
+    codeflash_return_value = codeflash_wrap(
+        collection.get_command,
+        "tests.test_sorter__unit_test_0",
+        None,
+        "test_command_exists_in_current_collection",
+        "get_command",
+        "2",
+        codeflash_loop_index,
+        mock_context,
+        "start",
+    )
+
+
+def test_command_exists_in_source_group(mock_context, mock_command, mock_group):
+    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
+    mock_group.get_command.return_value = mock_command
+    collection = CommandCollection(sources=[mock_group])
+    codeflash_return_value = codeflash_wrap(
+        collection.get_command,
+        "tests.test_sorter__unit_test_0",
+        None,
+        "test_command_exists_in_source_group",
+        "get_command",
+        "2",
+        codeflash_loop_index,
+        mock_context,
+        option="stop",
+    )
+
+
+def test_command_does_not_exist(mock_context, mock_group):
+    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
+    collection = CommandCollection(sources=[mock_group])
+    args = ["start", "stop", "restart"]
+    kwargs = {"option2": "stop"}
+    codeflash_return_value = codeflash_wrap(
+        collection.get_command,
+        "tests.test_sorter__unit_test_0",
+        None,
+        "test_command_does_not_exist",
+        "get_command",
+        "3",
+        codeflash_loop_index,
+        mock_context,
+        *args,
+        option="restart",
+        **kwargs,
+    )
+
+
+def test_empty_command_name(mock_context, mock_group):
+    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
+    collection = CommandCollection(sources=[mock_group])
+    codeflash_return_value = codeflash_wrap(
+        collection.get_command,
+        "tests.test_sorter__unit_test_0",
+        None,
+        "test_empty_command_name",
+        "get_command",
+        "1",
+        codeflash_loop_index,
+        mock_context,
+    )
+
+
+def test_none_as_command_name(mock_context, mock_group):
+    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
+    collection = CommandCollection(sources=[mock_group])
+    codeflash_return_value = codeflash_wrap(
+        collection.get_command,
+        "tests.test_sorter__unit_test_0",
+        None,
+        "test_none_as_command_name",
+        "get_command",
+        "1",
+        codeflash_loop_index,
+        mock_context,
+        sources=[mock_group],
+    )
+"""
+    )
+    function_to_optimize = FunctionToOptimize(
+        function_name="get_command",
+        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort_in_class.py",
+        parents=[FunctionParent(name="CommandCollection", type="ClassDef")],
+        starting_line=None,
+        ending_line=None,
+    )
+    result = instrument_test_source(
+        test_source=test_source,
+        function_to_optimize=function_to_optimize,
+        helper_function_names=[],
+        module_path="code_to_optimize.bubble_sort_in_class",
+        test_module_path="tests.test_sorter__unit_test_0",
+        test_framework="pytest",
+        test_timeout=15,
+        mode=TestingMode.PERFORMANCE,
+    )
+    assert result == expected
+
+
 def test_instrument_random_test() -> None:
     test_source = """import random
 def test_large_list_with_random_values():

django/aiservice/testgen/tests/test_perf_injector.py

@@ -45,7 +45,7 @@ def test_InjectPerfAndLogging_with() -> None:
     codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
     hdbscan = HDBSCAN()
     with pytest.raises(AttributeError):
-        codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1_0', codeflash_loop_index, )"""
+        codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1_0', codeflash_loop_index)"""
     assert ast.unparse(new_module_node).strip("\n") == expected
 
 
@@ -78,7 +78,7 @@ def test_InjectPerfAndLogging() -> None:
     expected = """def test_relative_validity_no_tree():
     codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
     hdbscan = HDBSCAN()
-    codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1', codeflash_loop_index, )"""
+    codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1', codeflash_loop_index)"""
     assert ast.unparse(new_module_node).strip("\n") == expected