diff --git a/local_code_assistant/commands/test.py b/local_code_assistant/commands/test.py new file mode 100644 index 0000000..2620f3f --- /dev/null +++ b/local_code_assistant/commands/test.py @@ -0,0 +1,210 @@ +"""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 + ) \ No newline at end of file