DeepSeek-OCR-2-Math / tests /test_example_loader.py
ricklon's picture
Handle more example payload shapes in workspace loader
2c2efb5
import ast
import os
import pathlib
import types
import unittest
def _load_example_loader():
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 = {
"_to_rgba_image",
"_scale_workspace_image",
"_prepare_workspace_image",
"load_example_into_workspace",
}
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")
class _FakeLoadedImage:
def __init__(self, width=1068, height=3074):
self.width = width
self.height = height
class _FakeImageModule:
class Image: # pragma: no cover - marker type for isinstance checks
pass
@staticmethod
def open(path):
return _FakeLoadedImage()
fake_np = types.SimpleNamespace(
ndarray=type("ndarray", (), {}),
uint8=int,
stack=lambda *args, **kwargs: None,
full_like=lambda *args, **kwargs: None,
concatenate=lambda *args, **kwargs: None,
)
scope = {
"os": os,
"np": fake_np,
"Image": _FakeImageModule,
"WORKSPACE_DEFAULT_SCALE": 89,
"load_image": lambda file_path, page_num=1: _FakeLoadedImage(),
}
exec(code, scope)
return scope["load_example_into_workspace"]
class ExampleLoaderTests(unittest.TestCase):
def test_accepts_common_gradio_payload_shapes(self):
loader = _load_example_loader()
sample = pathlib.Path("examples/2022-0922 Section 15 Notes.png")
for payload in (
str(sample),
sample,
{"path": str(sample)},
{"name": str(sample)},
[str(sample)],
(str(sample),),
):
display, size, base = loader(payload)
self.assertIsNotNone(display)
self.assertIsNotNone(base)
self.assertEqual((1068, 3074), size)
if __name__ == "__main__":
unittest.main()