chopratejas--headroom/tests/test_code_aware_brace_comment_regressions.py

"""Regression tests for two CodeAwareCompressor AST-reassembly bugs found
while investigating a reported Go brace-duplication issue.

1. `export` keyword duplication (TS/JS): `_compress_function_ast` and
   `_compress_class_ast` use LINE-based slicing (not byte-offset) to
   preserve indentation for nested definitions. When a function/class shares
   its first line with a preceding sibling — e.g. the `export` keyword in
   `export function foo() {`, a sibling of the function inside
   `export_statement`, not part of the function node itself — a naive
   full-line slice pulled that sibling's text in too. The `export_statement`
   handler then re-prepended the same `export` text, producing
   `export export function foo() {` (invalid syntax, silently discarded by
   the `_verify_syntax` fallback).
2. Doc-comment displacement (all languages): doc comments are top-level
   siblings of the declaration they document, not children of it. Left
   unattached during AST extraction, they fell through to a "leftover
   top-level code" bucket that `_assemble_compressed` emits as one block
   after every function signature — detaching every doc comment from what it
   documents and dumping them all at the end of the file.
"""

from __future__ import annotations

import pytest

from headroom.transforms.code_compressor import (
    CodeAwareCompressor,
    CodeCompressorConfig,
    CodeLanguage,
    _check_tree_sitter_available,
)

pytestmark = pytest.mark.skipif(
    not _check_tree_sitter_available(),
    reason="tree-sitter not installed (pip install headroom-ai[code])",
)

TS_EXPORTED = """export interface User {
  id: string;
  name: string;
}

/**
 * Fetches a user by id.
 */
export function getUser(id: string): User {
  return { id, name: "test" };
}

/**
 * Greets a user by name.
 */
export function greet(user: User): string {
  return `Hello, ${user.name}!`;
}

export class UserStore {
  private users: User[] = [];

  add(user: User): void {
    this.users.push(user);
  }
}
"""

GO_DOC_COMMENTS = """package main

import "fmt"

// Add adds two integers together and returns the sum.
func Add(a, b int) int {
\treturn a + b
}

// Greet returns a friendly greeting for the given name.
func Greet(name string) string {
\treturn fmt.Sprintf("Hello, %s!", name)
}
"""


def _compress_ast(code: str, language: CodeLanguage):
    compressor = CodeAwareCompressor(CodeCompressorConfig())
    compressed, _, _ = compressor._compress_with_ast(code, language, "", None)
    return compressor, compressed


def test_ts_export_keyword_not_duplicated() -> None:
    """`export function`/`export class` must not become `export export ...`."""
    compressor, compressed = _compress_ast(TS_EXPORTED, CodeLanguage.TYPESCRIPT)

    assert "export export" not in compressed, compressed
    assert compressor._verify_syntax(compressed, CodeLanguage.TYPESCRIPT) is True


def test_ts_doc_comments_stay_attached_to_declaration() -> None:
    """A `/** ... */` doc comment must stay immediately before the function it
    documents, not get dumped in a cluster at the end of the output."""
    _, compressed = _compress_ast(TS_EXPORTED, CodeLanguage.TYPESCRIPT)

    assert "/**\n * Fetches a user by id.\n */\nexport function getUser" in compressed
    assert "/**\n * Greets a user by name.\n */\nexport function greet" in compressed


def test_go_doc_comments_stay_attached_to_function() -> None:
    """Same doc-comment-attachment bug, Go's `//` line-comment form."""
    compressor, compressed = _compress_ast(GO_DOC_COMMENTS, CodeLanguage.GO)

    assert "// Add adds two integers together and returns the sum.\nfunc Add" in compressed
    assert "// Greet returns a friendly greeting for the given name.\nfunc Greet" in compressed
    assert compressor._verify_syntax(compressed, CodeLanguage.GO) is True


def test_actual_typescript_compression() -> None:
    """Parity with the existing JS/Python/Go 'actual compression' tests —
    real TS input must actually compress, not silently no-op."""
    config = CodeCompressorConfig(min_tokens_for_compression=10, enable_ccr=False)
    compressor = CodeAwareCompressor(config)
    code = TS_EXPORTED * 3  # large enough to trigger body elision

    result = compressor.compress(code, language="typescript")

    assert result.compression_ratio < 1.0
    assert result.syntax_valid is True
    assert result.language == CodeLanguage.TYPESCRIPT