Spaces:
Sleeping
Sleeping
| import sys | |
| import subprocess # nosec B404 | |
| from app.core.logger import get_logger | |
| from pathlib import Path | |
| from app.tools.file_scanner import get_python_files, sanitize_for_prompt | |
| from app.tools.ast_parser import parse_python_file | |
| from app.core.llm import call_llm_routed | |
| from app.core.prompts import TEST_GENERATION_PROMPT | |
| from app.models.review import GeneratedTests | |
| logger = get_logger(__name__) | |
| def _measure_coverage(local_path: str) -> float: | |
| try: | |
| subprocess.run( # nosec B603 B607 | |
| [ | |
| sys.executable, | |
| "-m", | |
| "coverage", | |
| "run", | |
| "--source=.", | |
| "-m", | |
| "pytest", | |
| "-q", | |
| "--no-header", | |
| ], | |
| cwd=local_path, | |
| capture_output=True, | |
| text=True, | |
| timeout=60, | |
| ) | |
| cov = subprocess.run( # nosec B603 B607 | |
| [sys.executable, "-m", "coverage", "report", "--format=total"], | |
| cwd=local_path, | |
| capture_output=True, | |
| text=True, | |
| timeout=30, | |
| ) | |
| return float(cov.stdout.strip()) if cov.returncode == 0 else 0.0 | |
| except Exception: | |
| return 0.0 | |
| async def run(local_path: str) -> GeneratedTests: | |
| """ | |
| Generate pytest test cases for uncovered functions. | |
| """ | |
| logger.info("Starting test generation agent", extra={"path": local_path}) | |
| # Step 1 - Get Python files | |
| python_files = get_python_files(local_path) | |
| # Step 2 - Find functions without tests | |
| all_functions = [] | |
| source_code_samples = [] | |
| for file_path in python_files: | |
| # Skip test files | |
| filename = Path(file_path).name | |
| if filename.startswith("test_") or filename.endswith("_test.py"): | |
| continue | |
| parsed = parse_python_file(file_path) | |
| functions = parsed.get("functions", []) | |
| if functions: | |
| all_functions.extend([f["name"] for f in functions]) | |
| try: | |
| with open(file_path, "r", encoding="utf-8", errors="ignore") as f: | |
| content = f.read()[:3000] | |
| content = sanitize_for_prompt(content) | |
| source_code_samples.append(f"### {file_path}\n{content}") | |
| except Exception: # nosec B110 | |
| pass | |
| if len(source_code_samples) >= 3: | |
| break | |
| if not source_code_samples: | |
| return GeneratedTests( | |
| test_code="# No source files found to generate tests for", | |
| total_tests_generated=0, | |
| ) | |
| # Step 3 - Call LLM | |
| source_code = "\n\n".join(source_code_samples) | |
| prompt = TEST_GENERATION_PROMPT.format(source_code=source_code) | |
| llm_response = await call_llm_routed(prompt, max_tokens=3000, task="test_generation") | |
| # Strip markdown code fences if present | |
| if llm_response.strip().startswith("```"): | |
| lines = llm_response.strip().splitlines() | |
| llm_response = "\n".join( | |
| line | |
| for i, line in enumerate(lines) | |
| if not (i == 0 and line.startswith("```")) and line != "```" | |
| ) | |
| # Step 4 - Count generated tests | |
| test_count = llm_response.count("def test_") | |
| coverage_before = _measure_coverage(local_path) | |
| result = GeneratedTests( | |
| test_code=llm_response, | |
| functions_covered=all_functions[:test_count], | |
| estimated_coverage_before=coverage_before, | |
| estimated_coverage_after=0.0, # not estimated; run generated tests to measure actual coverage | |
| total_tests_generated=test_count, | |
| ) | |
| logger.info("Test generation complete", extra={"tests_generated": test_count}) | |
| return result | |