"""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 )