markitect-main/tddai/test_generator.py

"""
Test generation with AI assistance.
"""

import subprocess
import tempfile
from pathlib import Path
from typing import Optional

from .config import get_config
from .workspace import WorkspaceManager
from .exceptions import TestGenerationError


class TestGenerator:
    """Generates tests using AI assistance."""

    def __init__(self, config=None):
        self.config = config or get_config()
        self.workspace_manager = WorkspaceManager(config)

    def generate_test(self, scenario_name: str, test_description: str) -> Path:
        """Generate a test file for the current workspace issue."""
        workspace = self.workspace_manager.get_current_workspace()
        if not workspace:
            raise TestGenerationError("No active workspace found")

        # Create test file name
        test_filename = self.config.test_file_pattern.format(
            issue_num=workspace.issue_number,
            scenario=scenario_name.lower().replace(' ', '_').replace('-', '_')
        )
        test_file_path = workspace.tests_dir / test_filename

        # Generate test prompt
        prompt = self._create_test_prompt(workspace, scenario_name, test_description)

        # Use Claude Code to generate the test
        try:
            with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
                f.write(prompt)
                prompt_file = Path(f.name)

            result = subprocess.run(
                [self.config.claude_code_command, '--file', str(prompt_file)],
                cwd=workspace.workspace_dir,
                capture_output=True,
                text=True
            )

            prompt_file.unlink()  # Clean up temp file

            if result.returncode != 0:
                raise TestGenerationError(f"Claude Code failed: {result.stderr}")

            # Extract Python code from Claude's response
            test_content = self._extract_test_code(result.stdout)

            # Write test file
            test_file_path.write_text(test_content)

            # Update test plan
            self._update_test_plan(workspace, scenario_name, test_filename)

            return test_file_path

        except subprocess.CalledProcessError as e:
            raise TestGenerationError(f"Failed to generate test: {e}")
        except Exception as e:
            raise TestGenerationError(f"Test generation error: {e}")

    def _create_test_prompt(self, workspace, scenario_name: str, test_description: str) -> str:
        """Create prompt for Claude Code to generate test."""
        return f"""# Test Generation Request

## Context
- Issue #{workspace.issue_number}: {workspace.issue_title}
- Scenario: {scenario_name}

## Issue Description
{workspace.issue_body}

## Test Requirements
{test_description}

## Instructions
Please generate a comprehensive Python test file that:

1. Tests the behavior described in the scenario
2. Follows pytest conventions
3. Includes proper docstrings and comments
4. Tests both positive and negative cases
5. Uses meaningful test method names
6. Includes appropriate assertions

The test should focus on behavior verification rather than implementation details.

## Expected Output
Please provide only the Python test code without any additional explanation.
The code should be ready to save as `{self.config.test_file_pattern.format(issue_num=workspace.issue_number, scenario=scenario_name.lower().replace(' ', '_'))}`
"""

    def _extract_test_code(self, claude_response: str) -> str:
        """Extract Python test code from Claude's response."""
        lines = claude_response.split('\n')
        code_lines = []
        in_code_block = False

        for line in lines:
            if line.strip().startswith('```python'):
                in_code_block = True
                continue
            elif line.strip() == '```' and in_code_block:
                break
            elif in_code_block:
                code_lines.append(line)

        if not code_lines:
            # If no code block found, assume entire response is code
            return claude_response.strip()

        return '\n'.join(code_lines)

    def _update_test_plan(self, workspace, scenario_name: str, test_filename: str) -> None:
        """Update the test plan with the new test."""
        test_plan_content = workspace.test_plan_file.read_text()

        # Add test to the generated tests section
        new_entry = f"- [x] {scenario_name} (`{test_filename}`)"

        if "### Generated Tests" in test_plan_content:
            # Add to existing generated tests section
            lines = test_plan_content.split('\n')
            for i, line in enumerate(lines):
                if line.strip() == "Tests generated for this workspace will be listed here as they are created.":
                    lines[i] = new_entry
                    break
                elif line.startswith("- [") and "Generated Tests" in lines[max(0, i-5):i]:
                    lines.insert(i, new_entry)
                    break
            else:
                # Add at the end of generated tests section
                for i, line in enumerate(lines):
                    if "### Generated Tests" in line:
                        # Find next section or end
                        j = i + 1
                        while j < len(lines) and not lines[j].startswith('##'):
                            j += 1
                        lines.insert(j, new_entry)
                        break

            workspace.test_plan_file.write_text('\n'.join(lines))

    def list_generated_tests(self) -> list:
        """List all generated tests for the current workspace."""
        workspace = self.workspace_manager.get_current_workspace()
        if not workspace:
            return []

        if not workspace.tests_dir.exists():
            return []

        return list(workspace.tests_dir.glob("*.py"))