"""Tests for postprocessing_testgen_pipeline."""

import libcst as cst

from aiservice.models.functions_to_optimize import FunctionToOptimize
from core.languages.python.testgen.postprocessing.postprocess_pipeline import postprocessing_testgen_pipeline


def test_replaces_function_definition_with_import() -> None:
    """LLMs sometimes include the full function definition instead of importing it.

    The postprocessing pipeline should replace the definition with an import statement.
    """
    generated_code = """\
import pytest

def comparator(orig, new, superset_obj=False):
    if type(orig) is not type(new):
        return False
    if isinstance(orig, (list, tuple)):
        if len(orig) != len(new):
            return False
        return all(comparator(e1, e2, superset_obj) for e1, e2 in zip(orig, new))
    if isinstance(orig, dict):
        if len(orig) != len(new):
            return False
        return all(k in new and comparator(v, new[k], superset_obj) for k, v in orig.items())
    return orig == new


def test_basic_equality():
    assert comparator(1, 1) is True
    assert comparator(1, 2) is False


def test_list_comparison():
    assert comparator([1, 2, 3], [1, 2, 3]) is True
    assert comparator([1, 2], [1, 2, 3]) is False
"""

    function_to_optimize = FunctionToOptimize(
        function_name="comparator", file_path="/path/to/codeflash/verification/comparator.py", parents=[]
    )
    source_code = "def comparator(orig, new):\n    return orig == new"

    result, _ = postprocessing_testgen_pipeline(
        module=cst.parse_module(generated_code),
        helper_function_names=[],
        function_to_optimize=function_to_optimize,
        module_path="codeflash.verification.comparator",
        source_code_being_tested=source_code,
    )

    output = result.code
    assert "from codeflash.verification.comparator import comparator" in output
    assert "def comparator(" not in output
    assert "def test_basic_equality" in output
    assert "def test_list_comparison" in output
