160 lines
5.0 KiB
Python
160 lines
5.0 KiB
Python
"""Tests for command generator."""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
|
|
|
|
class TestCommandGenerator:
|
|
"""Tests for CommandGenerator class."""
|
|
|
|
def test_generator_initialization(self):
|
|
"""Test CommandGenerator initialization."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
generator = CommandGenerator(mock_backend)
|
|
|
|
assert generator.backend == mock_backend
|
|
|
|
@patch('shellgen.core.generator.PromptTemplates')
|
|
def test_generate_basic_command(self, mock_templates):
|
|
"""Test generating a basic command."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
mock_backend.generate.return_value = "```\nls -la\n```\nThis lists all files with details."
|
|
|
|
generator = CommandGenerator(mock_backend)
|
|
mock_templates_instance = Mock()
|
|
mock_templates.generate_prompt.return_value = "Generate: list files"
|
|
mock_templates.return_value = mock_templates_instance
|
|
|
|
result = generator.generate("list files", shell="bash")
|
|
|
|
assert "ls" in result.command.lower()
|
|
assert len(result.explanation) > 0
|
|
|
|
def test_generate_zsh_shell(self):
|
|
"""Test generating command for Zsh."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
mock_backend.generate.return_value = "ls"
|
|
|
|
generator = CommandGenerator(mock_backend)
|
|
|
|
with patch.object(generator, 'templates') as mock_templates:
|
|
mock_templates.generate_prompt.return_value = "Generate for zsh"
|
|
result = generator.generate("list files", shell="zsh")
|
|
|
|
mock_templates.generate_prompt.assert_called_once()
|
|
call_args = mock_templates.generate_prompt.call_args
|
|
assert "zsh" in call_args[0][1].lower()
|
|
|
|
def test_extract_command_with_code_block(self):
|
|
"""Test extracting command from code block."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
mock_backend.generate.return_value = """Here is the command:
|
|
|
|
```
|
|
find . -name "*.py" -type f
|
|
```
|
|
|
|
This searches for Python files."""
|
|
|
|
generator = CommandGenerator(mock_backend)
|
|
result = generator.generate("find python files")
|
|
|
|
assert "find" in result.command
|
|
assert "name" in result.command or "*.py" in result.command
|
|
|
|
def test_extract_command_plain(self):
|
|
"""Test extracting command from plain text."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
mock_backend.generate.return_value = "git status"
|
|
|
|
generator = CommandGenerator(mock_backend)
|
|
result = generator.generate("show git status")
|
|
|
|
assert "git" in result.command
|
|
|
|
def test_confidence_estimation(self):
|
|
"""Test confidence estimation."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
mock_backend.generate.return_value = "```\nls\n```"
|
|
|
|
generator = CommandGenerator(mock_backend)
|
|
result = generator.generate("list files")
|
|
|
|
assert 0.0 <= result.confidence <= 1.0
|
|
|
|
def test_empty_response_fallback(self):
|
|
"""Test handling of empty response."""
|
|
from shellgen.core.generator import CommandGenerator
|
|
|
|
mock_backend = Mock()
|
|
mock_backend.generate.return_value = ""
|
|
|
|
generator = CommandGenerator(mock_backend)
|
|
result = generator.generate("list files")
|
|
|
|
assert result.command == ""
|
|
|
|
|
|
class TestPromptTemplates:
|
|
"""Tests for prompt templates."""
|
|
|
|
def test_generate_prompt_bash(self):
|
|
"""Test generating prompt for Bash."""
|
|
from shellgen.core.prompts import PromptTemplates
|
|
|
|
templates = PromptTemplates()
|
|
prompt = templates.generate_prompt("list files", "bash")
|
|
|
|
assert "BASH" in prompt
|
|
assert "list files" in prompt
|
|
|
|
def test_generate_prompt_zsh(self):
|
|
"""Test generating prompt for Zsh."""
|
|
from shellgen.core.prompts import PromptTemplates
|
|
|
|
templates = PromptTemplates()
|
|
prompt = templates.generate_prompt("find files", "zsh")
|
|
|
|
assert "ZSH" in prompt
|
|
assert "find files" in prompt
|
|
|
|
def test_generate_prompt_with_context(self):
|
|
"""Test generating prompt with context."""
|
|
from shellgen.core.prompts import PromptTemplates
|
|
|
|
templates = PromptTemplates()
|
|
prompt = templates.generate_prompt(
|
|
"list files",
|
|
"bash",
|
|
context="/home/user/project"
|
|
)
|
|
|
|
assert "/home/user/project" in prompt
|
|
|
|
def test_generate_feedback_prompt(self):
|
|
"""Test generating feedback prompt."""
|
|
from shellgen.core.prompts import PromptTemplates
|
|
|
|
templates = PromptTemplates()
|
|
prompt = templates.generate_feedback_prompt(
|
|
original_description="list files",
|
|
generated_command="ls -x",
|
|
user_correction="ls -la",
|
|
)
|
|
|
|
assert "list files" in prompt
|
|
assert "ls -x" in prompt
|
|
assert "ls -la" in prompt
|