tests/test_clean_output.py

import ast
import pathlib
import re
import unittest


def _load_clean_output():
    app_path = pathlib.Path(__file__).resolve().parents[1] / "app.py"
    source = app_path.read_text(encoding="utf-8")
    module = ast.parse(source, filename=str(app_path))

    wanted = {
        "_equation_text_key",
        "_dedupe_repeated_math_blocks",
        "_strip_malformed_grounding",
        "clean_output",
    }
    fn_nodes = [n for n in module.body if isinstance(n, ast.FunctionDef) and n.name in wanted]
    fn_nodes.sort(key=lambda n: n.lineno)

    test_mod = ast.Module(body=fn_nodes, type_ignores=[])
    code = compile(test_mod, filename=str(app_path), mode="exec")

    scope = {
        "re": re,
    }
    exec(code, scope)
    return scope["clean_output"]


class CleanOutputTests(unittest.TestCase):
    def test_removes_truncated_grounding_artifact_line(self):
        clean_output = _load_clean_output()
        raw = (
            "\\[ \\frac{18x-34}{(2x-3)^2} \\]\n"
            "<|ref|>equation<|/ref|><|det|>[[50, 0, 450, 100]]\n"
        )

        cleaned = clean_output(raw, include_images=True)
        self.assertIn("\\[ \\frac{18x-34}{(2x-3)^2} \\]", cleaned)
        self.assertNotIn("<|ref|>", cleaned)
        self.assertNotIn("<|det|>", cleaned)
        self.assertNotIn("[[50, 0, 450, 100]]", cleaned)

    def test_replaces_full_image_reference(self):
        clean_output = _load_clean_output()
        raw = "prefix\n<|ref|>image<|/ref|><|det|>[[0,0,100,100]]<|/det|>\nsuffix"

        cleaned = clean_output(raw, include_images=True)
        self.assertIn("**[Figure 1]**", cleaned)
        self.assertNotIn("<|ref|>", cleaned)
        self.assertNotIn("<|det|>", cleaned)

    def test_dedupes_equivalent_math_blocks(self):
        clean_output = _load_clean_output()
        raw = (
            "\\[ \\frac{(18x-27-7)}{(2x-3)^2}=\\frac{18x-34}{(2x-3)^2} \\]\n"
            "\\[ \\frac{(18x-27-7)}{(2x-3)^{2}}=\\frac{18x-34}{(2x-3)^{2}} \\]\n"
        )

        cleaned = clean_output(raw, include_images=True)
        self.assertEqual(1, cleaned.count("\\["))


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