src/test_generator.py

"""
AI Test Generator - Generates tests from requirements
"""

import os
from typing import List, Optional, TYPE_CHECKING
if TYPE_CHECKING:
    from src.types import PRDiff

from src.types import Requirement, GeneratedTest
from src.config import get_config, TestConfig
from src.ai_client import get_ai_client, AIClient


class TestGenerator:
    """Generates automated tests from requirements"""
    
    def __init__(self, ai_client: Optional[AIClient] = None, config: Optional[TestConfig] = None):
        self.config = config or get_config().tests
        self.ai = ai_client or get_ai_client()
    
    def generate_tests(
        self, 
        requirements: List[Requirement],
        pr_diff: Optional["PRDiff"] = None
    ) -> List[GeneratedTest]:
        """Generate tests for each requirement"""
        
        tests = []
        
        for requirement in requirements:
            if requirement.status.value == "PASS":
                # Generate test for passed requirements to verify
                test = self._generate_single_test(requirement)
                tests.append(test)
            else:
                # Generate test that SHOULD pass when requirement is implemented
                test = self._generate_single_test(requirement)
                test.error = "Requirement not yet implemented"
                tests.append(test)
        
        return tests
    
    def _generate_single_test(self, requirement: Requirement) -> GeneratedTest:
        """Generate a single test for a requirement"""
        
        system_prompt = f"""You are a test generation expert. Generate a {self.config.framework} test 
function that tests the given requirement. Return ONLY the test code, no explanations.
The test should be complete and runnable."""
        
        prompt = f"""Generate a test for this requirement:

Requirement ID: {requirement.id}
Text: {requirement.text}
Category: {requirement.category}
Priority: {requirement.priority}
Acceptance Criteria:
{chr(10).join(f"- {ac}" for ac in requirement.acceptance_criteria)}

Generate a {self.config.framework} test function. Include:
1. Proper imports
2. Clear test name (test_{requirement.id.lower()}_{requirement.category})
3. Assert statements for the acceptance criteria
4. Example test data if needed

Return ONLY the test code, no markdown formatting."""
        
        try:
            test_code = self.ai.complete(prompt, system_prompt)
            
            # Clean up the test code
            test_code = test_code.strip()
            if test_code.startswith("```python"):
                test_code = test_code[10:]
            if test_code.startswith("```"):
                test_code = test_code[3:]
            if test_code.endswith("```"):
                test_code = test_code[:-3]
            test_code = test_code.strip()
            
            return GeneratedTest(
                name=f"test_{requirement.id.lower()}_{requirement.category}",
                requirement_id=requirement.id,
                test_code=test_code,
                framework=self.config.framework,
            )
            
        except Exception as e:
            return GeneratedTest(
                name=f"test_{requirement.id.lower()}_{requirement.category}",
                requirement_id=requirement.id,
                test_code=f"# Error generating test: {str(e)}",
                framework=self.config.framework,
                error=str(e),
            )
    
    def save_tests(self, tests: List[GeneratedTest]) -> List[str]:
        """Save generated tests to files"""
        
        os.makedirs(self.config.output_dir, exist_ok=True)
        
        saved_files = []
        
        # Group tests by requirement ID
        tests_by_req = {}
        for test in tests:
            if test.requirement_id not in tests_by_req:
                tests_by_req[test.requirement_id] = []
            tests_by_req[test.requirement_id].append(test)
        
        # Save each group
        for req_id, req_tests in tests_by_req.items():
            filename = os.path.join(self.config.output_dir, f"test_{req_id.lower()}.py")
            
            # Build test file content
            content = f'''"""
Auto-generated tests for requirement {req_id}
Generated by AI PR Reviewer
"""
import pytest


'''
            
            for test in req_tests:
                content += test.test_code + "\n\n"
            
            with open(filename, 'w') as f:
                f.write(content)
            
            saved_files.append(filename)
        
        return saved_files
    
    def run_tests(self, tests: List[GeneratedTest]) -> List[GeneratedTest]:
        """Run generated tests (if auto_run is enabled)"""
        
        if not self.config.auto_run:
            return tests
        
        # Save tests first
        self.save_tests(tests)
        
        # Try to run pytest
        try:
            import subprocess
            result = subprocess.run(
                ["pytest", self.config.output_dir, "-v"],
                capture_output=True,
                text=True,
            )
            
            # Update test results
            for test in tests:
                test.passed = result.returncode == 0
                if result.returncode != 0:
                    test.error = result.stderr
            
            return tests
            
        except Exception as e:
            for test in tests:
                test.error = f"Failed to run tests: {str(e)}"
            return tests