claude-skills-reference/engineering-team/tdd-guide/test_generator.py

"""
Test case generation module.

Generates test cases from requirements, user stories, API specs, and code analysis.
Supports multiple testing frameworks with intelligent test scaffolding.
"""

from typing import Dict, List, Any, Optional
from enum import Enum


class TestFramework(Enum):
    """Supported testing frameworks."""
    JEST = "jest"
    VITEST = "vitest"
    PYTEST = "pytest"
    JUNIT = "junit"
    MOCHA = "mocha"


class TestType(Enum):
    """Types of tests to generate."""
    UNIT = "unit"
    INTEGRATION = "integration"
    E2E = "e2e"


class TestGenerator:
    """Generate test cases and test stubs from requirements and code."""

    def __init__(self, framework: TestFramework, language: str):
        """
        Initialize test generator.

        Args:
            framework: Testing framework to use
            language: Programming language (typescript, javascript, python, java)
        """
        self.framework = framework
        self.language = language
        self.test_cases = []

    def generate_from_requirements(
        self,
        requirements: Dict[str, Any],
        test_type: TestType = TestType.UNIT
    ) -> List[Dict[str, Any]]:
        """
        Generate test cases from requirements.

        Args:
            requirements: Dictionary with user_stories, acceptance_criteria, api_specs
            test_type: Type of tests to generate

        Returns:
            List of test case specifications
        """
        test_cases = []

        # Generate from user stories
        if 'user_stories' in requirements:
            for story in requirements['user_stories']:
                test_cases.extend(self._test_cases_from_story(story))

        # Generate from acceptance criteria
        if 'acceptance_criteria' in requirements:
            for criterion in requirements['acceptance_criteria']:
                test_cases.extend(self._test_cases_from_criteria(criterion))

        # Generate from API specs
        if 'api_specs' in requirements:
            for endpoint in requirements['api_specs']:
                test_cases.extend(self._test_cases_from_api(endpoint))

        self.test_cases = test_cases
        return test_cases

    def _test_cases_from_story(self, story: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Generate test cases from user story."""
        test_cases = []

        # Happy path test
        test_cases.append({
            'name': f"should_{story.get('action', 'work')}_successfully",
            'type': 'happy_path',
            'description': story.get('description', ''),
            'given': story.get('given', []),
            'when': story.get('when', ''),
            'then': story.get('then', ''),
            'priority': 'P0'
        })

        # Error cases
        if 'error_conditions' in story:
            for error in story['error_conditions']:
                test_cases.append({
                    'name': f"should_handle_{error.get('condition', 'error')}",
                    'type': 'error_case',
                    'description': error.get('description', ''),
                    'expected_error': error.get('error_type', ''),
                    'priority': 'P0'
                })

        # Edge cases
        if 'edge_cases' in story:
            for edge_case in story['edge_cases']:
                test_cases.append({
                    'name': f"should_handle_{edge_case.get('scenario', 'edge_case')}",
                    'type': 'edge_case',
                    'description': edge_case.get('description', ''),
                    'priority': 'P1'
                })

        return test_cases

    def _test_cases_from_criteria(self, criterion: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Generate test cases from acceptance criteria."""
        return [{
            'name': f"should_meet_{criterion.get('id', 'criterion')}",
            'type': 'acceptance',
            'description': criterion.get('description', ''),
            'verification': criterion.get('verification_steps', []),
            'priority': 'P0'
        }]

    def _test_cases_from_api(self, endpoint: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Generate test cases from API specification."""
        test_cases = []
        method = endpoint.get('method', 'GET')
        path = endpoint.get('path', '/')

        # Success case
        test_cases.append({
            'name': f"should_{method.lower()}_{path.replace('/', '_')}_successfully",
            'type': 'api_success',
            'method': method,
            'path': path,
            'expected_status': endpoint.get('success_status', 200),
            'priority': 'P0'
        })

        # Validation errors
        if 'required_params' in endpoint:
            test_cases.append({
                'name': f"should_return_400_for_missing_params",
                'type': 'api_validation',
                'method': method,
                'path': path,
                'expected_status': 400,
                'priority': 'P0'
            })

        # Authorization
        if endpoint.get('requires_auth', False):
            test_cases.append({
                'name': f"should_return_401_for_unauthenticated",
                'type': 'api_auth',
                'method': method,
                'path': path,
                'expected_status': 401,
                'priority': 'P0'
            })

        return test_cases

    def generate_test_stub(self, test_case: Dict[str, Any]) -> str:
        """
        Generate test stub code for a test case.

        Args:
            test_case: Test case specification

        Returns:
            Test stub code as string
        """
        if self.framework == TestFramework.JEST:
            return self._generate_jest_stub(test_case)
        elif self.framework == TestFramework.PYTEST:
            return self._generate_pytest_stub(test_case)
        elif self.framework == TestFramework.JUNIT:
            return self._generate_junit_stub(test_case)
        elif self.framework == TestFramework.VITEST:
            return self._generate_vitest_stub(test_case)
        else:
            return self._generate_generic_stub(test_case)

    def _generate_jest_stub(self, test_case: Dict[str, Any]) -> str:
        """Generate Jest test stub."""
        name = test_case.get('name', 'test')
        description = test_case.get('description', '')

        stub = f"""
describe('{{Feature Name}}', () => {{
  it('{name}', () => {{
    // {description}

    // Arrange
    // TODO: Set up test data and dependencies

    // Act
    // TODO: Execute the code under test

    // Assert
    // TODO: Verify expected behavior
    expect(true).toBe(true); // Replace with actual assertion
  }});
}});
"""
        return stub.strip()

    def _generate_pytest_stub(self, test_case: Dict[str, Any]) -> str:
        """Generate Pytest test stub."""
        name = test_case.get('name', 'test')
        description = test_case.get('description', '')

        stub = f"""
def test_{name}():
    \"\"\"
    {description}
    \"\"\"
    # Arrange
    # TODO: Set up test data and dependencies

    # Act
    # TODO: Execute the code under test

    # Assert
    # TODO: Verify expected behavior
    assert True  # Replace with actual assertion
"""
        return stub.strip()

    def _generate_junit_stub(self, test_case: Dict[str, Any]) -> str:
        """Generate JUnit test stub."""
        name = test_case.get('name', 'test')
        description = test_case.get('description', '')

        # Convert snake_case to camelCase for Java
        method_name = ''.join(word.capitalize() if i > 0 else word
                             for i, word in enumerate(name.split('_')))

        stub = f"""
@Test
public void {method_name}() {{
    // {description}

    // Arrange
    // TODO: Set up test data and dependencies

    // Act
    // TODO: Execute the code under test

    // Assert
    // TODO: Verify expected behavior
    assertTrue(true); // Replace with actual assertion
}}
"""
        return stub.strip()

    def _generate_vitest_stub(self, test_case: Dict[str, Any]) -> str:
        """Generate Vitest test stub (similar to Jest)."""
        name = test_case.get('name', 'test')
        description = test_case.get('description', '')

        stub = f"""
describe('{{Feature Name}}', () => {{
  it('{name}', () => {{
    // {description}

    // Arrange
    // TODO: Set up test data and dependencies

    // Act
    // TODO: Execute the code under test

    // Assert
    // TODO: Verify expected behavior
    expect(true).toBe(true); // Replace with actual assertion
  }});
}});
"""
        return stub.strip()

    def _generate_generic_stub(self, test_case: Dict[str, Any]) -> str:
        """Generate generic test stub."""
        name = test_case.get('name', 'test')
        description = test_case.get('description', '')

        return f"""
# Test: {name}
# Description: {description}
#
# TODO: Implement test
# 1. Arrange: Set up test data
# 2. Act: Execute code under test
# 3. Assert: Verify expected behavior
"""

    def generate_test_file(
        self,
        module_name: str,
        test_cases: Optional[List[Dict[str, Any]]] = None
    ) -> str:
        """
        Generate complete test file with all test stubs.

        Args:
            module_name: Name of module being tested
            test_cases: List of test cases (uses self.test_cases if not provided)

        Returns:
            Complete test file content
        """
        cases = test_cases or self.test_cases

        if self.framework == TestFramework.JEST:
            return self._generate_jest_file(module_name, cases)
        elif self.framework == TestFramework.PYTEST:
            return self._generate_pytest_file(module_name, cases)
        elif self.framework == TestFramework.JUNIT:
            return self._generate_junit_file(module_name, cases)
        elif self.framework == TestFramework.VITEST:
            return self._generate_vitest_file(module_name, cases)
        else:
            return ""

    def _generate_jest_file(self, module_name: str, test_cases: List[Dict[str, Any]]) -> str:
        """Generate complete Jest test file."""
        imports = f"import {{ {module_name} }} from '../{module_name}';\n\n"

        stubs = []
        for test_case in test_cases:
            stubs.append(self._generate_jest_stub(test_case))

        return imports + "\n\n".join(stubs)

    def _generate_pytest_file(self, module_name: str, test_cases: List[Dict[str, Any]]) -> str:
        """Generate complete Pytest test file."""
        imports = f"import pytest\nfrom {module_name} import *\n\n\n"

        stubs = []
        for test_case in test_cases:
            stubs.append(self._generate_pytest_stub(test_case))

        return imports + "\n\n\n".join(stubs)

    def _generate_junit_file(self, module_name: str, test_cases: List[Dict[str, Any]]) -> str:
        """Generate complete JUnit test file."""
        class_name = ''.join(word.capitalize() for word in module_name.split('_'))

        imports = """import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.*;

"""

        class_header = f"public class {class_name}Test {{\n\n"

        stubs = []
        for test_case in test_cases:
            stubs.append(self._generate_junit_stub(test_case))

        class_footer = "\n}"

        return imports + class_header + "\n\n".join(stubs) + class_footer

    def _generate_vitest_file(self, module_name: str, test_cases: List[Dict[str, Any]]) -> str:
        """Generate complete Vitest test file."""
        imports = f"import {{ describe, it, expect }} from 'vitest';\nimport {{ {module_name} }} from '../{module_name}';\n\n"

        stubs = []
        for test_case in test_cases:
            stubs.append(self._generate_vitest_stub(test_case))

        return imports + "\n\n".join(stubs)

    def suggest_missing_scenarios(
        self,
        existing_tests: List[str],
        code_analysis: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """
        Suggest missing test scenarios based on code analysis.

        Args:
            existing_tests: List of existing test names
            code_analysis: Analysis of code under test (branches, error paths, etc.)

        Returns:
            List of suggested test scenarios
        """
        suggestions = []

        # Check for untested error conditions
        if 'error_handlers' in code_analysis:
            for error_handler in code_analysis['error_handlers']:
                error_name = error_handler.get('type', 'error')
                if not self._has_test_for(existing_tests, error_name):
                    suggestions.append({
                        'name': f"should_handle_{error_name}",
                        'type': 'error_case',
                        'reason': 'Error handler exists but no corresponding test',
                        'priority': 'P0'
                    })

        # Check for untested branches
        if 'conditional_branches' in code_analysis:
            for branch in code_analysis['conditional_branches']:
                branch_name = branch.get('condition', 'condition')
                if not self._has_test_for(existing_tests, branch_name):
                    suggestions.append({
                        'name': f"should_test_{branch_name}_branch",
                        'type': 'branch_coverage',
                        'reason': 'Conditional branch not fully tested',
                        'priority': 'P1'
                    })

        # Check for boundary conditions
        if 'input_validation' in code_analysis:
            for validation in code_analysis['input_validation']:
                param = validation.get('parameter', 'input')
                if not self._has_test_for(existing_tests, f"{param}_boundary"):
                    suggestions.append({
                        'name': f"should_test_{param}_boundary_values",
                        'type': 'boundary',
                        'reason': 'Input validation exists but boundary tests missing',
                        'priority': 'P1'
                    })

        return suggestions

    def _has_test_for(self, existing_tests: List[str], keyword: str) -> bool:
        """Check if existing tests cover a keyword/scenario."""
        keyword_lower = keyword.lower().replace('_', '').replace('-', '')
        for test in existing_tests:
            test_lower = test.lower().replace('_', '').replace('-', '')
            if keyword_lower in test_lower:
                return True
        return False