"""Tests for parse_and_validate_llm_output function."""

from aiservice.models.functions_to_optimize import FunctionToOptimize
from core.languages.python.testgen.generate import parse_and_validate_llm_output


def test_preserves_test_code_when_no_function_to_replace() -> None:
    """Test that code is preserved when there's no function definition to replace."""
    llm_response = """```python
import pytest
from mymodule import my_function


def test_basic():
    assert my_function(1) == 2
```
"""

    function_to_optimize = FunctionToOptimize(function_name="my_function", file_path="/path/to/mymodule.py", parents=[])
    module_path = "mymodule"

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        source_code="def my_function(x):\n    return x + 1",
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=function_to_optimize,
        module_path=module_path,
    )

    expected = """import pytest
from mymodule import my_function


def test_basic():
    assert my_function(1) == 2"""

    assert result.strip() == expected.strip()


def test_works_without_function_to_optimize() -> None:
    """Test that validation still works when function_to_optimize is None."""
    llm_response = """```python
import pytest


def test_simple():
    assert 1 == 1
```
"""

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        source_code="def foo():\n    return 1",
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=None,
        module_path=None,
    )

    expected = """import pytest


def test_simple():
    assert 1 == 1"""

    assert result.strip() == expected.strip()
