tests/test_spatial_blocks.py

import ast
import pathlib
import re
import unittest


def _load_grounding_blocks():
    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 = {
        "_parse_coord_payload",
        "_extract_grounding_entries",
        "_grounding_blocks_from_raw",
    }
    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 = {
        "ast": ast,
        "re": re,
        "GROUNDING_PATTERN": re.compile(r"<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>", re.DOTALL),
    }
    exec(code, scope)
    return scope["_grounding_blocks_from_raw"]


class SpatialBlockTests(unittest.TestCase):
    def test_multi_coord_refs_render_as_multiple_blocks(self):
        grounding_blocks = _load_grounding_blocks()
        raw = (
            "Equation cluster\n"
            "<|ref|>formula<|/ref|><|det|>[[100,100,600,220],[100,240,600,360]]<|/det|>\n"
            "Trailing text\n"
            "<|ref|>text<|/ref|><|det|>[[40,400,700,520]]<|/det|>"
        )

        blocks = grounding_blocks(raw)
        self.assertEqual(3, len(blocks))

        formula_blocks = [b for b in blocks if b["label"] == "formula"]
        self.assertEqual(2, len(formula_blocks))
        self.assertEqual(100.0, formula_blocks[0]["x1"])
        self.assertEqual(600.0, formula_blocks[0]["x2"])
        self.assertEqual(100.0, formula_blocks[1]["x1"])
        self.assertEqual(600.0, formula_blocks[1]["x2"])
        self.assertEqual("Equation cluster", formula_blocks[0]["text"])
        self.assertEqual("", formula_blocks[1]["text"])


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