import pytest from unittest.mock import Mock, patch, MagicMock from pathlib import Path import json import os import sys from datetime import datetime from src.llm_codegen.core import CodeGenerator from src.llm_codegen.models import GeneratorState, ConfigModel from src.llm_codegen.utils import is_dangerous_command class TestCodeGenerator: """测试 CodeGenerator 核心类的单元测试。""" @pytest.fixture def mock_openai_client(self): """模拟 OpenAI 客户端。""" with patch('src.llm_codegen.core.OpenAI') as mock: client = Mock() mock.return_value = client yield client @pytest.fixture def generator(self, mock_openai_client, tmp_path): """创建 CodeGenerator 实例，使用临时目录和模拟 API。""" output_dir = tmp_path / "output" return CodeGenerator( api_key="test-api-key", base_url="https://api.deepseek.com", model="deepseek-reasoner", output_dir=str(output_dir), resume=False, ) def test_init(self, generator, tmp_path): """测试初始化。""" assert generator.api_key == "test-api-key" assert generator.model == "deepseek-reasoner" assert generator.output_dir == tmp_path / "output" assert generator.resume is False assert isinstance(generator.config, ConfigModel) assert generator.dangerous_commands == ["rm", "sudo", "chmod", "dd", "mkfs", "> /dev/sda", "format"] assert generator.state is None def test_parse_readme(self, generator, tmp_path): """测试读取 README 文件。""" readme_path = tmp_path / "README.md" readme_content = "# Test Project\nThis is a test README." readme_path.write_text(readme_content) result = generator.parse_readme(readme_path) assert result == readme_content def test_get_project_structure(self, generator, mock_openai_client): """测试获取项目结构，模拟 LLM 响应。""" generator.readme_content = "# Test README" mock_response = { "files": ["src/__init__.py", "src/core.py"], "dependencies": {"src/core.py": ["src/__init__.py"]} } mock_openai_client.chat.completions.create.return_value.choices[0].message.content = json.dumps(mock_response) files, dependencies = generator.get_project_structure() assert files == ["src/__init__.py", "src/core.py"] assert dependencies == {"src/core.py": ["src/__init__.py"]} mock_openai_client.chat.completions.create.assert_called_once() def test_generate_file(self, generator, mock_openai_client, tmp_path): """测试生成单个文件，模拟依赖文件和 LLM 响应。""" generator.readme_content = "# Test README" dep_file = tmp_path / "dep.txt" dep_file.write_text("Dependency content") mock_response = { "code": "print('Hello, World!')", "description": "测试文件生成", "commands": ["echo 'test'"] } mock_openai_client.chat.completions.create.return_value.choices[0].message.content = json.dumps(mock_response) code, desc, commands = generator.generate_file( "test.py", "生成测试文件", [str(dep_file)] ) assert code == "print('Hello, World!')" assert desc == "测试文件生成" assert commands == ["echo 'test'"] mock_openai_client.chat.completions.create.assert_called_once() def test_execute_command_safe(self, generator, tmp_path): """测试执行安全命令。""" with patch('subprocess.run') as mock_run: mock_run.return_value.returncode = 0 mock_run.return_value.stdout = "output" mock_run.return_value.stderr = "" generator.execute_command("echo 'test'", cwd=tmp_path) mock_run.assert_called_once_with( "echo 'test'", shell=True, cwd=tmp_path, capture_output=True, text=True, timeout=300 ) def test_execute_command_dangerous(self, generator): """测试阻止危险命令。""" with pytest.raises(RuntimeError, match="危险命令"): generator.execute_command("rm -rf /") # 假设在危险命令列表中 def test_run_without_resume(self, generator, mock_openai_client, tmp_path): """测试完整运行流程，禁用断点续写。""" readme_path = tmp_path / "README.md" readme_path.write_text("# Test README") generator.readme_content = "# Test README" # 模拟 get_project_structure 响应 mock_structure = { "files": ["file1.py", "file2.py"], "dependencies": {} } mock_openai_client.chat.completions.create.side_effect = [ Mock(choices=[Mock(message=Mock(content=json.dumps(mock_structure)))]), Mock(choices=[Mock(message=Mock(content=json.dumps({"code": "code1", "description": "desc1", "commands": []})))]), Mock(choices=[Mock(message=Mock(content=json.dumps({"code": "code2", "description": "desc2", "commands": []})))]) ] with patch('src.llm_codegen.core.safe_write_file') as mock_write, \ patch('src.llm_codegen.core.safe_read_file') as mock_read, \ patch('src.llm_codegen.core.save_state') as mock_save: mock_read.return_value = "content" generator.run(readme_path) # 验证文件生成和状态保存 assert mock_write.call_count == 2 assert mock_save.called def test_run_with_resume(self, generator, mock_openai_client, tmp_path): """测试断点续写功能。""" generator.resume = True generator.state = GeneratorState(generated_files=["file1.py"], executed_commands=[]) readme_path = tmp_path / "README.md" readme_path.write_text("# Test README") generator.readme_content = "# Test README" mock_structure = { "files": ["file1.py", "file2.py"], "dependencies": {} } mock_openai_client.chat.completions.create.return_value.choices[0].message.content = json.dumps(mock_structure) with patch('src.llm_codegen.core.safe_write_file') as mock_write, \ patch('src.llm_codegen.core.safe_read_file') as mock_read: mock_read.return_value = "content" generator.run(readme_path) # 只应生成 file2.py，跳过 file1.py assert mock_write.call_count == 1 def test_load_config_default(self, generator): """测试加载默认配置。""" config = generator._load_config() assert isinstance(config, ConfigModel) assert config.check_tools == ["pytest", "pylint", "mypy", "black"] assert config.max_retries == 3 def test_update_state(self, generator, tmp_path): """测试更新状态文件。""" generator.state_path = tmp_path / "state.json" generator.state = GeneratorState() with patch('src.llm_codegen.core.save_state') as mock_save: generator._update_state("new_file.py", ["cmd1"]) assert generator.state.generated_files == ["new_file.py"] assert generator.state.executed_commands == ["cmd1"] mock_save.assert_called_once() if __name__ == "__main__": pytest.main([__file__, "-v"])