compounding-intelligence/scripts/test_documentation_generator.py

#!/usr/bin/env python3
"""Test Documentation Generator — adds module and function docstrings to test files.

Reads test files without docstrings and generates:
  - Module-level docstring explaining what is being tested
  - Function-level docstring explaining what each test verifies
  - Inline comments for complex assertions (simple heuristic)

Does not change test logic — only adds documentation.
Processes 20+ test files per run.
"""

import ast
import re
import sys
from pathlib import Path
from typing import List, Tuple


def derive_module_name(test_path: Path) -> str:
    """Derive the script/module name being tested from test file name."""
    name = test_path.stem
    if name.startswith("test_"):
        name = name[5:]  # strip 'test_'  (5 chars: t-e-s-t-_, not 6)
    mapping = {
        "bootstrapper": "bootstrapper.py",
        "harvester": "harvester.py",
        "diff_analyzer": "diff_analyzer.py",
        "gitea_issue_parser": "gitea_issue_parser.py",
        "harvest_prompt": "harvest_prompt.py",
        "harvest_prompt_comprehensive": "harvest_prompt_comprehensive.py",
        "harvester_pipeline": "harvester_pipeline.py",
        "improvement_proposals": "improvement_proposals.py",
        "knowledge_staleness": "knowledge_staleness_check.py",
        "priority_rebalancer": "priority_rebalancer.py",
        "refactoring_opportunity_finder": "refactoring_opportunity_finder.py",
        "session_pair_harvester": "session_pair_harvester.py",
        "session_reader": "session_reader.py",
        "automation_opportunity_finder": "automation_opportunity_finder.py",
        "dedup": "dedup.py",
        "freshness": "freshness.py",
        "knowledge_gap_identifier": "knowledge_gap_identifier.py",
        "perf_bottleneck_finder": "perf_bottleneck_finder.py",
        "ci_config": "CI configuration",
        "quality_gate": "quality_gate.py",
    }
    base = name.replace("_", " ")
    if name in mapping:
        base = mapping[name].replace(".py", "")
    return base


def count_tests_in_file(content: str) -> int:
    """Count test functions in a Python file."""
    return len(re.findall(r'^def (test_\w+)\s*\(', content, re.MULTILINE))


def infer_test_purpose(func_name: str, func_body: str) -> str:
    """Generate a brief docstring for a test function based on its name and body."""
    name = func_name.replace("test_", "").replace("_", " ")

    if "empty" in name or "none" in name:
        return "Verifies behavior with empty or None input."
    if "parsing" in name or "parse" in name:
        return f"Verifies parsing logic for {name}."
    if "filter" in name:
        return f"Verifies knowledge filtering by {name}."
    if "hash" in name:
        return "Verifies file hash computation correctness."
    if "freshness" in name or "staleness" in name:
        return "Verifies knowledge freshness detection."
    if "error" in name or "exception" in name:
        return f"Verifies error handling for {name}."
    if "boundary" in name or "edge" in name:
        return "Verifies boundary case handling."
    return f"Verifies {name} logic."


def has_module_docstring(content: str) -> bool:
    """Check if file (after shebang) starts with a docstring."""
    lines = content.split('\n')
    start_idx = 1 if lines and lines[0].startswith('#!') else 0
    for line in lines[start_idx:start_idx + 5]:
        stripped = line.strip()
        if stripped.startswith('"""') or stripped.startswith("'''"):
            return True
        if stripped == "" or stripped.startswith('#'):
            continue
        break
    return False


def insert_after_shebang(content: str, insertion: str) -> str:
    """Insert text after the shebang line (if any) and any following blank lines."""
    lines = content.split('\n')
    insert_idx = 0
    if lines and lines[0].startswith('#!'):
        insert_idx = 1
        while insert_idx < len(lines) and lines[insert_idx].strip() == '':
            insert_idx += 1
    new_lines = lines[:insert_idx] + [insertion] + lines[insert_idx:]
    return '\n'.join(new_lines)


def add_function_docstring(content: str, func_lineno: int, docstring: str) -> str:
    """Add a docstring to a function at the given line number."""
    lines = content.split('\n')
    idx = func_lineno - 1
    indent = re.match(r'^(\s*)', lines[idx]).group(1)
    doc_line = f'{indent}    """{docstring}"""'
    new_lines = lines[:idx + 1] + [doc_line] + lines[idx + 1:]
    return '\n'.join(new_lines)


def generate_module_docstring(test_path: Path) -> str:
    """Generate a module-level docstring for a test file."""
    module = derive_module_name(test_path)
    count = count_tests_in_file(test_path.read_text())
    if count > 0:
        return f"Tests for {module} — {count} tests."
    return f"Tests for {module}."


def process_test_file(test_path: Path, dry_run: bool = False) -> Tuple[bool, List[str]]:
    """Process a single test file, adding missing docstrings. Returns (changed, messages)."""
    content = test_path.read_text()
    original = content
    messages = []

    if not has_module_docstring(content):
        mod_doc = generate_module_docstring(test_path)
        content = insert_after_shebang(content, f'''"""{mod_doc}"""''')
        messages.append(f"Added module docstring: {mod_doc}")

    try:
        tree = ast.parse(content)
    except SyntaxError as e:
        messages.append(f"SKIP (syntax error): {e}")
        return False, messages

    funcs_to_doc: List[Tuple[int, str, str]] = []

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name.startswith('test_'):
            has_docstring = (
                len(node.body) > 0 and
                isinstance(node.body[0], ast.Expr) and
                isinstance(node.body[0].value, ast.Constant) and
                isinstance(node.body[0].value.value, str)
            )
            if not has_docstring:
                func_body = ast.get_source_segment(content, node) or ""
                doc = infer_test_purpose(node.name, func_body)
                funcs_to_doc.append((node.lineno, node.name, doc))

    funcs_to_doc.sort(key=lambda x: -x[0])
    for lineno, func_name, doc in funcs_to_doc:
        content = add_function_docstring(content, lineno, doc)
        messages.append(f"Added docstring to {func_name}: {doc}")

    changed = content != original
    if changed and not dry_run:
        test_path.write_text(content)

    return changed, messages


def find_test_files(root: Path, max_files: int = 25) -> List[Path]:
    """Find test files under scripts/ and tests/ directories."""
    test_files = []
    for subdir in [root / "scripts", root / "tests"]:
        if subdir.exists():
            test_files.extend(subdir.glob("test_*.py"))
    test_files.sort()
    return test_files[:max_files]


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Generate documentation for test files")
    parser.add_argument("--dry-run", action="store_true", help="Show changes without writing")
    parser.add_argument("--root", type=Path, default=Path.cwd(),
                        help="Repo root (default: current directory)")
    parser.add_argument("--limit", type=int, default=25,
                        help="Max files to process per run (handles 20+ requirement)")
    args = parser.parse_args()

    root = args.root
    test_files = find_test_files(root, args.limit)
    print(f"Found {len(test_files)} test files to process (limit={args.limit}):")

    total_changed = 0
    for tf in test_files:
        changed, msgs = process_test_file(tf, dry_run=args.dry_run)
        if changed:
            total_changed += 1
        status = "CHANGED" if changed else "OK"
        print(f"  [{status}] {tf.relative_to(root)}")
        for msg in msgs:
            print(f"       {msg}")

    print(f"\nCompleted: {total_changed} file(s) modified, {len(test_files) - total_changed} already up-to-date.")
    return 0


if __name__ == "__main__":
    sys.exit(main())