""" 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