DocForge2/tests/test_treesitter_logic.py

from __future__ import annotations

import unittest

from docforge.errors import DocForgeError
from docforge.treesitter_logic import (
    TreeSitterLogicOwner,
    analyze_cpp_source,
    analyze_javascript_source,
    analyze_typescript_source,
    discover_typescript_functions,
)


class JavaScriptLogicTests(unittest.TestCase):
    def test_consecutive_comments_become_one_logic_comment_block(self) -> None:
        source = """
function choose(enabled) {
  // Explain the condition.
  // Continue the explanation.
  /* A block comment is trivia too. */
  if (enabled) {
    accept();
  }
}
""".strip()
        projection = analyze_javascript_source(
            source,
            source_id="source.javascript",
            owners=(TreeSitterLogicOwner("js.symbol.choose", "choose", 1),),
        )[0]

        comments = [node for node in projection.nodes if node.kind == "comment"]
        self.assertEqual(len(comments), 1)
        self.assertEqual(
            comments[0].label,
            "Explain the condition. Continue the explanation. A block comment is trivia too.",
        )
        self.assertIn("accept();", {node.label for node in projection.nodes})

    def test_branches_short_circuit_and_converge(self) -> None:
        source = """
function choose(enabled, ready) {
  if (enabled && ready()) {
    accept();
  } else {
    reject();
  }
  return enabled;
}
""".strip()
        projection = analyze_javascript_source(
            source,
            source_id="source.javascript",
            owners=(TreeSitterLogicOwner("js.symbol.choose", "choose", 1),),
        )[0]

        kinds = {node.kind for node in projection.nodes}
        labels = {node.label for node in projection.nodes}
        relations = {edge.relation for edge in projection.edges}
        self.assertIn("condition", kinds)
        self.assertIn("convergence", kinds)
        self.assertIn("Decision convergence", labels)
        self.assertIn("when_true", relations)
        self.assertIn("when_false", relations)
        self.assertIn("return", relations)

    def test_arrow_function_expression_is_a_returning_projection(self) -> None:
        source = "const choose = (enabled) => enabled ? accept() : reject();"
        projection = analyze_javascript_source(
            source,
            source_id="source.javascript",
            owners=(TreeSitterLogicOwner("js.symbol.choose", "choose", 1),),
        )[0]

        labels = {node.label for node in projection.nodes}
        self.assertIn("enabled", labels)
        self.assertIn("return accept()", labels)
        self.assertIn("return reject()", labels)

    def test_loops_switch_and_exception_paths_are_preserved(self) -> None:
        source = """
function process(items, mode) {
  for (const item of items) {
    if (!item.ready) continue;
    use(item);
  }
  switch (mode) {
    case 1:
      one();
      break;
    default:
      fallback();
  }
  try {
    risk();
  } catch (error) {
    recover(error);
  } finally {
    clean();
  }
}
""".strip()
        projection = analyze_javascript_source(
            source,
            source_id="source.javascript",
            owners=(TreeSitterLogicOwner("js.symbol.process", "process", 1),),
        )[0]

        kinds = {node.kind for node in projection.nodes}
        labels = {node.label for node in projection.nodes}
        relations = {edge.relation for edge in projection.edges}
        self.assertTrue({"loop", "continue", "case", "try", "except", "finally"} <= kinds)
        self.assertIn("Case convergence", labels)
        self.assertIn("Exception convergence", labels)
        self.assertTrue({"loop", "continue", "case", "exception"} <= relations)


class CppLogicTests(unittest.TestCase):
    def test_consecutive_comments_become_one_logic_comment_block(self) -> None:
        source = """
int choose(bool enabled) {
  // Explain the condition.
  /* A block comment is trivia too. */
  if (enabled) {
    return 1;
  }
  return 0;
}
""".strip()
        projection = analyze_cpp_source(
            source,
            source_id="source.cpp",
            owners=(TreeSitterLogicOwner("cpp.symbol.choose", "choose", 1),),
        )[0]

        comments = [node for node in projection.nodes if node.kind == "comment"]
        self.assertEqual(len(comments), 1)
        self.assertEqual(
            comments[0].label,
            "Explain the condition. A block comment is trivia too.",
        )
        self.assertIn("return 1", {node.label for node in projection.nodes})

    def test_cpp_function_branches_and_throws(self) -> None:
        source = """
int choose(bool enabled) {
  if (enabled) {
    return 1;
  }
  throw Error();
}
""".strip()
        projection = analyze_cpp_source(
            source,
            source_id="source.cpp",
            owners=(TreeSitterLogicOwner("cpp.symbol.choose", "choose", 1),),
        )[0]

        kinds = {node.kind for node in projection.nodes}
        relations = {edge.relation for edge in projection.edges}
        self.assertTrue({"entry", "condition", "return", "raise", "exit"} <= kinds)
        self.assertTrue({"when_true", "when_false", "return", "raise"} <= relations)

    def test_cpp_qualified_method_owner_is_resolved(self) -> None:
        source = """
class Worker {
public:
  int run(bool ready) {
    while (ready) {
      ready = tick();
    }
    return 0;
  }
};
""".strip()
        projection = analyze_cpp_source(
            source,
            source_id="source.cpp",
            owners=(TreeSitterLogicOwner("cpp.symbol.worker.run", "Worker.run", 3),),
        )[0]

        labels = {node.label for node in projection.nodes}
        self.assertIn("Loop exit", labels)
        self.assertIn("return 0", labels)

    def test_invalid_source_and_missing_owner_fail_closed(self) -> None:
        with self.assertRaises(DocForgeError) as invalid:
            analyze_cpp_source(
                "int broken( {",
                source_id="source.cpp",
                owners=(),
            )
        self.assertEqual(invalid.exception.code, "invalid_logic_source")

        with self.assertRaises(DocForgeError) as missing:
            analyze_javascript_source(
                "function exists() {}",
                source_id="source.javascript",
                owners=(TreeSitterLogicOwner("js.symbol.missing", "missing", 1),),
            )
        self.assertEqual(missing.exception.code, "missing_logic_owner")


class TypeScriptLogicTests(unittest.TestCase):
    def test_typed_function_and_method_use_the_typescript_grammar(self) -> None:
        source = """
interface Choice {
  enabled: boolean;
}

function choose(choice: Choice): number {
  if (choice.enabled) {
    return 1;
  }
  return 0;
}

class Worker {
  run(choice: Choice): number {
    return choose(choice);
  }
}
""".strip()

        discovered = discover_typescript_functions(source)
        self.assertEqual(
            [("Worker.run", 13), ("choose", 5)],
            [(item.qualified_name, item.line) for item in discovered],
        )
        projection = analyze_typescript_source(
            source,
            source_id="source.typescript",
            owners=(TreeSitterLogicOwner("ts.symbol.choose", "choose", 5),),
        )[0]

        self.assertIn("choice.enabled", {node.label for node in projection.nodes})
        self.assertEqual(
            {"return"},
            {edge.relation for edge in projection.edges if edge.relation == "return"},
        )


if __name__ == "__main__":
    unittest.main()