local-code-assistant/local_code_assistant/commands/test.py

"""Test command for Local Code Assistant."""

from pathlib import Path
from typing import Optional

import click
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax

from local_code_assistant.commands.base import BaseCommand
from local_code_assistant.prompts.templates import LanguageConfig, PromptTemplates
from local_code_assistant.services.ollama import OllamaService

console = Console()


class TestCommand(BaseCommand):
    """Command for generating tests."""

    def __init__(self, ollama: OllamaService, config):
        """Initialize test command.

        Args:
            ollama: Ollama service instance.
            config: Configuration service instance.
        """
        super().__init__(ollama, config)
        self.supported_languages = LanguageConfig.get_supported_languages()

    def run(
        self,
        code: str,
        language: str,
        framework: Optional[str] = None,
        output: Optional[Path] = None,
        clipboard: bool = False
    ) -> str:
        """Execute test generation.

        Args:
            code: Code to generate tests for.
            language: Programming language.
            framework: Testing framework (auto-detected if not provided).
            output: Output file path.
            clipboard: Copy to clipboard.

        Returns:
            Generated test code.
        """
        if not self.ollama.check_connection():
            raise click.ClickException(
                "Cannot connect to Ollama. Make sure it's running."
            )

        if language not in self.supported_languages:
            raise click.ClickException(f"Unsupported language: {language}")

        console.print("[dim]Generating tests...[/dim]")

        full_prompt = PromptTemplates.test_generation(
            language=language,
            code=code
        )

        system_prompt = PromptTemplates.build_system_prompt(
            "You are writing tests. Be thorough and cover edge cases."
        )

        try:
            test_code = self.ollama.generate(
                prompt=full_prompt,
                model=self._get_model(),
                system=system_prompt,
                temperature=self._get_temperature(0.3)
            )

            self._display_output(test_code, language)

            if output:
                output.write_text(test_code)
                console.print(f"[green]Tests written to {output}[/green]")

            if clipboard:
                import pyperclip
                pyperclip.copy(test_code)
                console.print("[green]Tests copied to clipboard[/green]")

            return test_code

        except Exception as e:
            raise click.ClickException(f"Test generation failed: {str(e)}") from e

    def run_file(
        self,
        file: Path,
        framework: Optional[str] = None,
        output: Optional[Path] = None,
        clipboard: bool = False
    ) -> str:
        """Generate tests for a file.

        Args:
            file: Path to file.
            framework: Testing framework.
            output: Output file path.
            clipboard: Copy to clipboard.

        Returns:
            Generated test code.
        """
        if not file.exists():
            raise click.ClickException(f"File not found: {file}")

        code = file.read_text()
        language = self._detect_language_from_file(file)

        return self.run(
            code=code,
            language=language,
            framework=framework,
            output=output,
            clipboard=clipboard
        )

    def _detect_language_from_file(self, file: Path) -> str:
        """Detect language from file extension.

        Args:
            file: File path.

        Returns:
            Language name.
        """
        suffix = file.suffix.lower()

        language_map = {
            ".py": "python",
            ".js": "javascript",
            ".ts": "typescript",
            ".tsx": "typescript",
            ".go": "go",
            ".rs": "rust"
        }

        return language_map.get(suffix, "python")

    def _display_output(self, code: str, language: str):
        """Display generated test code.

        Args:
            code: Test code.
            language: Programming language.
        """
        if self.config.syntax_highlighting:
            syntax = Syntax(code, language, line_numbers=True)
            console.print(Panel(syntax, title="Generated Tests"))
        else:
            console.print(Panel(code, title="Generated Tests"))


@click.command()
@click.argument("file", type=click.Path(exists=True, path_type=Path))
@click.option(
    "--framework", "-f",
    help="Testing framework (auto-detected if not specified)"
)
@click.option(
    "--output", "-o",
    type=click.Path(path_type=Path),
    help="Write tests to file"
)
@click.option(
    "--clipboard/--no-clipboard",
    default=False,
    help="Copy tests to clipboard"
)
@click.pass_context
def test_cmd(
    ctx: click.Context,
    file: Path,
    framework: Optional[str],
    output: Optional[Path],
    clipboard: bool
):
    """Generate tests for a code file.

    Example:
        local-code-assistant test my_module.py
        local-code-assistant test app.py -o test_app.py

    \f
    Args:
        ctx: Click context.
        file: Path to file to generate tests for.
        framework: Testing framework.
        output: Output file path.
        clipboard: Copy to clipboard.
    """
    config = ctx.obj["config"]
    ollama_service = OllamaService(config)
    ctx.obj["ollama_service"] = ollama_service

    command = TestCommand(ollama_service, config)
    command.run_file(
        file=file,
        framework=framework,
        output=output,
        clipboard=clipboard
    )