agents/test_generator.py

"""
agents/test_generator.py
------------------------
Test Generator Agent for AutoDevAgent.

Responsible for producing a test suite for the generated code after
it has executed successfully. This closes the loop — code that runs
is not necessarily correct, but code that passes tests is much more
likely to be.

Design:
  - Python tasks: generates standard unittest test cases that import
    and call the generated code directly. Tests cover the happy path,
    edge cases, and type correctness.
  - SQL tasks: generates Python assertions that run additional queries
    against the same SQLite schema to verify row counts, column
    presence, value ranges, and null checks.
  - Uses the primary model (70B) — test quality matters as much as
    code quality.
  - Returns raw test code as a string — the TestRunner executes it
    separately so failures re-engage the debug loop.

Usage:
    from agents.test_generator import TestGeneratorAgent
    from pipeline.state import PipelineState, Language

    agent = TestGeneratorAgent()
    # state.generated_code and state.execution_result must be set
    updated = agent.run(state)
    print(updated["generated_tests"])
"""

import logging
from typing import Any

from langchain_groq import ChatGroq
from langchain_core.messages import SystemMessage, HumanMessage

from config import settings
from pipeline.state import (
    PipelineState,
    PipelineStatus,
    Language,
    SQLSchema,
)

logger = logging.getLogger(__name__)


# ------------------------------------------------------------------ #
#  Prompts                                                             #
# ------------------------------------------------------------------ #

PYTHON_TEST_SYSTEM = """
You are an expert Python test engineer.
You will be given a task description, the working code, and its output.
Write a unittest test suite that verifies the code is correct.

Rules:
- Output ONLY raw Python test code. No markdown, no explanation, no code fences.
- Use the standard unittest module — no pytest, no third-party libs.
- The test file must be self-contained and runnable as-is with: python -m unittest test_file.py
- Copy the function(s) being tested directly into the test file so it is standalone.
  Do NOT use import statements to import from external files.
- Write at least 3 test cases covering:
    1. A standard happy-path input
    2. An edge case (empty string, zero, single element, etc.)
    3. A boundary or negative test (e.g. inputs that should return False or 0)
- IMPORTANT — Do NOT copy examples from the code: The "Output when run" section shows one
  example execution. Do NOT simply copy those same input values into your test assertions.
  Your test inputs must be DIFFERENT values that you independently derive from the function's
  logic. Re-testing the exact same example adds no value and does not verify correctness.
  Example: if the code prints "fibonacci(10) = 55", do NOT write assertEqual(fibonacci(10), 55).
  Instead write assertEqual(fibonacci(7), 13) or assertEqual(fibonacci(0), 0) etc.
- IMPORTANT — None inputs: ONLY test None inputs if the task description explicitly mentions
  handling None or null values. For string/list tasks that don't mention None, do NOT write
  tests that pass None — they will cause TypeError errors in functions that don't handle it.
  Example: "check if two strings are anagrams" → do NOT test are_anagrams(None, "hello")
  Example: "handle None gracefully and return False" → DO test func(None)
- Every test method name must start with test_
- IMPORTANT: Every test method MUST have a one-sentence docstring that explains exactly
  what input is used and what the expected outcome is. Be specific — mention the actual
  values, not just "tests the happy path". Example:
      def test_empty_list(self):
          \"\"\"Verifies that passing an empty list returns 0 because there are no elements to sum.\"\"\"
- End the file with:
    if __name__ == "__main__":
        unittest.main(verbosity=2)
""".strip()

SQL_TEST_SYSTEM = """
You are an expert SQL test engineer working with SQLite.
You will be given:
  - The original SQL task
  - The SQL query that was written
  - The database schema (CREATE TABLE + dummy data)
  - The query output

Write Python assertion tests that verify the query is correct.
The tests will run the query again against the same SQLite schema and assert on the results.

Rules:
- Output ONLY raw Python code. No markdown, no explanation, no code fences.
- Use Python's built-in sqlite3 and unittest modules only.
- Re-create the schema and dummy data inside the test setUp() method.
- Always access result values by POSITION (row[0], row[1]) NOT by column name string.
  SQLite returns aggregate expressions as raw strings like "MAX(salary)" or "COUNT(*)"
  so column-name access is fragile and must never be used.
- Write exactly 3 assertion-based tests covering:
    1. Row count — the query returns the expected number of rows
    2. Not-empty check — the result set is non-empty and the first row exists
    3. Value check — at least one specific value in the results is correct,
       accessed by POSITION (row[0], row[1], etc.)
- IMPORTANT: Every test method MUST have a one-sentence docstring that explains exactly
  what is being verified and why the expected value is correct given the dummy data.
  Be specific — mention the actual expected values. Example:
      def test_row_count(self):
          \"\"\"Verifies the query returns exactly 2 rows because only Alice and Bob have more than 3 orders in the dummy data.\"\"\"
- The test file must be runnable as-is.
- End the file with:
    if __name__ == "__main__":
        unittest.main(verbosity=2)
""".strip()


# ------------------------------------------------------------------ #
#  Agent                                                               #
# ------------------------------------------------------------------ #

class TestGeneratorAgent:
    """
    Generates a test suite for successfully executed code.

    For Python: produces unittest test cases that call the generated
    functions directly. For SQL: produces Python assertion checks that
    re-run the query against the same SQLite schema.

    Attributes:
        llm: ChatGroq instance using the primary model (70B).
    """

    def __init__(self) -> None:
        """LLM client is built lazily in run() once model assignments are known."""
        self.llm = None

    def _build_llm(self, model: str) -> "ChatGroq":
        return ChatGroq(
            api_key=settings.groq_api_key,
            model=model,
            temperature=0.1,
            max_tokens=2000,
            request_timeout=settings.groq_request_timeout,
        )

    def run(self, state: PipelineState) -> dict[str, Any]:
        """
        Generate tests for the code in the current pipeline state.

        Builds a prompt containing the task, the working code, and
        its output. Sends to the LLM and returns the raw test code.

        Args:
            state: Current PipelineState. Reads: task, language,
                   generated_code, execution_result, sql_schema.

        Returns:
            Partial state dict with keys:
              - "generated_tests": str  — raw test code
              - "status":          PipelineStatus.TESTING
        """
        ma = state.model_assignments or {}
        model = ma.get("tester", settings.groq_model_primary)
        self.llm = self._build_llm(model)
        logger.info("TestGeneratorAgent using model: %s", model)

        logger.info(
            "TestGeneratorAgent running for language: %s", state.language.value
        )

        system_prompt = (
            SQL_TEST_SYSTEM
            if state.language == Language.SQL
            else PYTHON_TEST_SYSTEM
        )

        human_content = _build_human_message(state)

        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=human_content),
        ]

        # ── LLM call ──────────────────────────────────────────────── #
        try:
            response = self.llm.invoke(messages)
            raw      = response.content.strip()
            logger.debug(
                "TestGeneratorAgent raw response length: %d chars", len(raw)
            )
        except Exception as e:
            logger.error("TestGeneratorAgent LLM call failed: %s", e)
            raise RuntimeError(f"Test generator LLM call failed: {e}") from e

        # ── Strip markdown fences if LLM added them ───────────────── #
        tests = _strip_code_fences(raw)

        if not tests.strip():
            raise ValueError("TestGeneratorAgent returned empty test code.")

        logger.info(
            "TestGeneratorAgent produced %d lines of test code",
            len(tests.splitlines()),
        )

        # ── Track token usage ─────────────────────────────────────── #
        from observability.langsmith_tracer import extract_token_usage_from_response
        prompt_t, completion_t = extract_token_usage_from_response(response)
        updated_token_usage = state.token_usage.model_copy()
        updated_token_usage.add(prompt_t, completion_t)

        return {
            "generated_tests": tests,
            "status":          PipelineStatus.TESTING,
            "token_usage":     updated_token_usage,
        }


# ------------------------------------------------------------------ #
#  Helpers                                                             #
# ------------------------------------------------------------------ #

def _build_human_message(state: PipelineState) -> str:
    """
    Build the human message for the test generator LLM call.

    Includes the task, the working code, the execution output, and
    for SQL tasks the schema context so the LLM can recreate the DB.

    Args:
        state: Current PipelineState.

    Returns:
        Formatted prompt string.
    """
    stdout = ""
    if state.execution_result and state.execution_result.stdout:
        stdout = state.execution_result.stdout[:500]   # Cap to avoid prompt bloat

    lines: list[str] = [
        f"Task: {state.task}",
        "",
        "Working code:",
        state.generated_code,
        "",
        f"Output when run:",
        stdout if stdout else "(no stdout — code ran without printing)",
    ]

    # For SQL tasks include the schema so the LLM can recreate it in tests
    if state.language == Language.SQL and state.sql_schema:
        lines.append("")
        lines.append("Database schema:")
        for stmt in state.sql_schema.create_statements:
            lines.append(f"  {stmt}")
        lines.append("")
        lines.append("Dummy data:")
        for stmt in state.sql_schema.insert_statements:
            lines.append(f"  {stmt}")

    lines.append("")
    lines.append("Write the test suite now:")
    return "\n".join(lines)


def _strip_code_fences(raw: str) -> str:
    """
    Remove markdown code fences from LLM output.

    Args:
        raw: Raw string from the LLM response.

    Returns:
        Clean code string with fences removed.
    """
    stripped = raw.strip()
    if stripped.startswith("```"):
        lines = stripped.splitlines()
        start = 1
        end   = len(lines) - 1 if lines[-1].strip() == "```" else len(lines)
        return "\n".join(lines[start:end]).strip()
    return stripped