AI-система автоматической генерации unit-тестов по коду

AI-автогенерация unit-тестов

Unit-тесты пишут все, но мало кто делает это систематически — особенно для унаследованного кода, где нет тестов и непонятно с чего начинать. AI-генератор создаёт pytest/Jest/JUnit тесты непосредственно из исходного кода, анализируя логику через AST и выявляя сценарии, которые человек пропустил бы.

Python: генерация pytest через AST + LLM

import ast
import inspect
from langchain_openai import ChatOpenAI
from pathlib import Path

class UnitTestGenerator:
    PYTEST_PROMPT = """Сгенерируй pytest unit-тесты для функции.

Код функции:
```python
{function_code}

Зависимости модуля: {imports}

Анализ через AST:

• Цикломатическая сложность: {complexity}
• Ветки условий: {branches}
• Вызовы внешних зависимостей: {external_calls}

Требования к тестам:

1. Используй @pytest.mark.parametrize для наборов данных
2. Мок внешние зависимости через pytest-mock (mocker.patch)
3. Тестируй все ветки: каждое условие if/elif/else
4. Тестируй raises: для каждого raise в коде
5. Используй fixtures для переиспользуемых объектов
6. Имена тестов: test_{function_name}_{scenario} (напр. test_calculate_tax_zero_income)

Верни только код тестов с import-секцией."""

def __init__(self):
    self.llm = ChatOpenAI(model="gpt-4o", temperature=0.1)

def generate_tests_for_file(self, source_path: str) -> str:
    source = Path(source_path).read_text(encoding="utf-8")
    tree = ast.parse(source)

    all_tests = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if node.name.startswith("_"):
                continue  # пропускаем приватные методы

            func_source = ast.get_source_segment(source, node)
            analysis = self._analyze_function(node, source)

            tests = self._generate_function_tests(func_source, analysis, source)
            all_tests.append(tests)

    return self._merge_test_files(all_tests, source_path)

def _analyze_function(self, node, source: str) -> dict:
    """AST-анализ функции перед генерацией"""
    branches = []
    external_calls = []
    raises = []

    for child in ast.walk(node):
        if isinstance(child, ast.If):
            cond = ast.get_source_segment(source, child.test)
            branches.append(cond)
        elif isinstance(child, ast.Call):
            if isinstance(child.func, ast.Attribute):
                call = f"{ast.get_source_segment(source, child.func.value)}.{child.func.attr}"
                external_calls.append(call)
        elif isinstance(child, ast.Raise):
            if child.exc:
                raises.append(ast.get_source_segment(source, child.exc))

    return {
        "complexity": self._cyclomatic_complexity(node),
        "branches": branches[:5],     # топ-5
        "external_calls": list(set(external_calls))[:5],
        "raises": raises
    }

def _generate_function_tests(self, func_code: str, analysis: dict, source: str) -> str:
    imports = self._extract_imports(source)

    result = self.llm.invoke(
        self.PYTEST_PROMPT.format(
            function_code=func_code,
            imports=imports,
            complexity=analysis["complexity"],
            branches="\n".join(analysis["branches"]),
            external_calls="\n".join(analysis["external_calls"])
        )
    )
    return result.content

### TypeScript/Jest генерация

```python
    JEST_PROMPT = """Сгенерируй Jest unit-тесты для TypeScript функции.

```typescript
{function_code}

Требования:

• Используй describe/it блоки
• jest.fn() для моков
• beforeEach для setup
• expect().toBe() / toEqual() / toThrow()
• Тест должен импортировать только нужное
• Не используй any — только строгая типизация в тестах

Покрой: успешные сценарии, граничные значения, ошибки. Верни только TypeScript код."""

async def generate_jest_tests(self, ts_function: str) -> str:
    result = await self.llm.ainvoke(
        self.JEST_PROMPT.format(function_code=ts_function)
    )
    return result.content

### Автоматическое выявление проблем в сгенерированных тестах

Сгенерированные тесты иногда имеют синтаксические ошибки или неверные ассерты. Добавляем validation loop:

```python
import subprocess

class TestValidator:
    def validate_and_fix(self, test_code: str, source_file: str) -> str:
        """Запускает тесты и фиксит ошибки в цикле"""
        temp_test_file = "/tmp/test_generated.py"
        Path(temp_test_file).write_text(test_code)

        for attempt in range(3):
            result = subprocess.run(
                ["pytest", temp_test_file, "-x", "--tb=short",
                 f"--rootdir={Path(source_file).parent}"],
                capture_output=True, text=True, timeout=60
            )

            if result.returncode == 0:
                break

            # Фиксим ошибки через LLM
            fix_prompt = f"""Исправь pytest тесты, которые не проходят.

Тесты:
{test_code}

Ошибка:
{result.stdout[-2000:]}

Верни исправленные тесты (только код)."""

            test_code = self.llm.invoke(fix_prompt).content
            Path(temp_test_file).write_text(test_code)

        return test_code

Интеграция через pre-commit hook

# .pre-commit-config.yaml
repos:
  - repo: local
    hooks:
      - id: ai-test-generator
        name: AI Unit Test Generator
        entry: python scripts/generate_tests.py
        language: python
        pass_filenames: true
        types: [python]
        stages: [push]  # только при push, не при каждом commit

Кейс: Python-сервис обработки платежей, 12 000 строк кода, 0 unit-тестов (legacy). Запустили генератор на всю кодовую базу: 340 тестов за 45 минут. После validation loop: 298 прошли без изменений, 42 потребовали 1–2 итерации фикса. Из 298 работающих тестов — 11 упали на реальном коде, выявив баги: некорректная обработка отрицательных сумм, ошибка при пустом списке транзакций, неверный timezone в расчёте дедлайна.

Сроки: генератор для одного языка (Python/TypeScript) с validation loop: 2–3 недели; мультиязычный с CI/CD интеграцией: 4–5 недель.