164 lines
7.1 KiB
Python
164 lines
7.1 KiB
Python
"""Offline worker tests with fake inference; no model or source network access."""
|
|
|
|
import importlib.util
|
|
import math
|
|
import socket
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "clinical_image_ocr.py"
|
|
SPEC = importlib.util.spec_from_file_location("clinical_image_ocr", SCRIPT)
|
|
ocr = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(ocr)
|
|
|
|
|
|
class Image:
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
def detected(text="SYNTHETIC", confidence=0.9, box=None):
|
|
return SimpleNamespace(txts=[text], scores=[confidence], boxes=[box or [[0, 0], [1, 0], [1, 1], [0, 1]]])
|
|
|
|
|
|
class WorkerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.private = tempfile.TemporaryDirectory(prefix="ocr-worker-test-")
|
|
self.root = Path(self.private.name)
|
|
self.image = self.root / "image.png"
|
|
self.image.write_bytes(b"synthetic fake image for injected decoder")
|
|
self.request = {"allowed_root": str(self.root), "images": [
|
|
{"index": 7, "path": str(self.image)}, {"index": 19, "path": str(self.image)},
|
|
], "deadline_ms": 1000, "threads": 1, "detection_limit": 640, "max_text_bytes": 4096}
|
|
|
|
def tearDown(self):
|
|
self.private.cleanup()
|
|
|
|
def run_worker(self, request=None, factory=None, loader=None, checkpoint=None):
|
|
return ocr.run(request or self.request, engine_factory=factory or (lambda *_: lambda _image: detected()),
|
|
image_loader=loader or (lambda _: Image()), checkpoint=checkpoint)
|
|
|
|
def test_complete_global_indices_and_one_engine(self):
|
|
calls = []
|
|
|
|
def factory(threads, detection):
|
|
calls.append((threads, detection))
|
|
return lambda _image: detected()
|
|
|
|
result = self.run_worker(factory=factory)
|
|
self.assertEqual(result["status"], "completed")
|
|
self.assertEqual([row["index"] for row in result["images"]], [7, 19])
|
|
self.assertEqual(calls, [(1, 640)])
|
|
|
|
def test_invalid_indices_limits_paths_and_no_model_load(self):
|
|
variants = [
|
|
{"images": [{"index": 7, "path": str(self.image)}] * 2},
|
|
{"images": [{"index": True, "path": str(self.image)}]},
|
|
{"images": [{"index": 7, "path": str(SCRIPT)}]},
|
|
{"images": [{"index": 7, "path": "https://synthetic.invalid/image.png"}]},
|
|
{"images": [{"index": 7, "path": "\\\\synthetic.invalid\\share\\image.png"}]},
|
|
{"allowed_root": "//synthetic.invalid/share"},
|
|
{"threads": 3}, {"detection_limit": 10000}, {"max_text_bytes": 9999999},
|
|
{"deadline_ms": -1}, {"deadline_ms": 0},
|
|
]
|
|
for changed in variants:
|
|
with self.subTest(changed=list(changed)):
|
|
result = self.run_worker({**self.request, **changed}, factory=lambda *_: self.fail("must not load model"))
|
|
self.assertEqual(result["status"], "unavailable")
|
|
|
|
def test_too_many_bytes_and_symlink_outside(self):
|
|
with self.image.open("wb") as handle:
|
|
handle.truncate(ocr.MAX_BYTES + 1)
|
|
self.assertEqual(self.run_worker()["status"], "unavailable")
|
|
self.image.write_bytes(b"synthetic")
|
|
link = self.root / "outside.png"
|
|
try:
|
|
link.symlink_to(SCRIPT)
|
|
except OSError:
|
|
return
|
|
result = self.run_worker({**self.request, "images": [{"index": 7, "path": str(link)}]})
|
|
self.assertEqual(result["status"], "unavailable")
|
|
|
|
def test_pixel_header_guard_before_large_decode(self):
|
|
from PIL import Image as PillowImage
|
|
|
|
image = self.root / "large.bmp"
|
|
image.write_bytes(b"BM" + struct.pack("<IHHI", 54, 0, 0, 54)
|
|
+ struct.pack("<IiiHHIIiiII", 40, 10000, 10000, 1, 24, 0, 0, 0, 0, 0, 0))
|
|
with self.assertRaises((ocr.Unavailable, PillowImage.DecompressionBombError, PillowImage.DecompressionBombWarning)):
|
|
ocr.load_image(image)
|
|
|
|
def test_network_denied_and_socket_functions_restored(self):
|
|
original = socket.socket.connect
|
|
|
|
def factory(*_):
|
|
with self.assertRaises(OSError):
|
|
socket.create_connection(("example.test", 443))
|
|
return lambda _image: detected()
|
|
|
|
self.assertEqual(self.run_worker(factory=factory)["status"], "completed")
|
|
self.assertIs(socket.socket.connect, original)
|
|
|
|
def test_first_unreadable_does_not_block_other_source(self):
|
|
calls = []
|
|
|
|
def loader(_):
|
|
calls.append(True)
|
|
if len(calls) == 1:
|
|
raise OSError("private decoder error")
|
|
return Image()
|
|
|
|
result = self.run_worker(loader=loader)
|
|
self.assertEqual(result["status"], "partial")
|
|
self.assertEqual([row["status"] for row in result["images"]], ["unreadable", "processed"])
|
|
self.assertNotIn("private decoder error", str(result))
|
|
|
|
def test_deadline_checkpoint_preserves_completed_and_unfinished_rows(self):
|
|
clock, checkpoints = [0.0], []
|
|
|
|
def checkpoint(value):
|
|
checkpoints.append([row["status"] for row in value["images"]])
|
|
clock[0] = 2.0
|
|
|
|
with patch.object(ocr.time, "monotonic", side_effect=lambda: clock[0]):
|
|
result = self.run_worker(checkpoint=checkpoint)
|
|
self.assertEqual(result["status"], "partial")
|
|
self.assertEqual(result["reason"], "deadline")
|
|
self.assertEqual(checkpoints, [["processed", "unreadable"]])
|
|
|
|
def test_invalid_confidence_boxes_and_text_never_become_processed(self):
|
|
for output in [detected(confidence=math.nan), detected(confidence=2), detected(text="bad\0text"),
|
|
detected(box=[[math.inf, 0]] * 4), detected(text="x" * 9000)]:
|
|
with self.subTest():
|
|
result = self.run_worker(factory=lambda *_, value=output: lambda _image: value)
|
|
self.assertEqual(result["status"], "unavailable")
|
|
self.assertTrue(all(not row["lines"] for row in result["images"]))
|
|
|
|
def test_shared_text_budget_keeps_first_source_and_marks_later_unreadable(self):
|
|
result = self.run_worker({**self.request, "max_text_bytes": len("SYNTHETIC")})
|
|
self.assertEqual(result["status"], "partial")
|
|
self.assertEqual([row["status"] for row in result["images"]], ["processed", "unreadable"])
|
|
|
|
def test_dependency_failure_is_safe_and_covers_every_index(self):
|
|
def factory(*_):
|
|
raise RuntimeError("private model path and stderr")
|
|
|
|
result = self.run_worker(factory=factory)
|
|
self.assertEqual(result["status"], "unavailable")
|
|
self.assertEqual(len(result["images"]), 2)
|
|
self.assertNotIn("private", str(result))
|
|
|
|
def test_missing_local_models_fail_before_import_or_download(self):
|
|
with patch.object(ocr.importlib.metadata, "version", side_effect=lambda name: {"rapidocr": "3.10.0", "onnxruntime": "1.31.0"}[name]), \
|
|
patch.object(ocr.importlib.util, "find_spec", return_value=SimpleNamespace(origin=str(self.root / "__init__.py"))), \
|
|
self.assertRaises(FileNotFoundError):
|
|
ocr.load_engine(1, 640)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|