"""OCR 流水线：多模态识别、Pix2Text 降级、提取失败兜底。""" import io import json from unittest.mock import patch import pytest from PIL import Image as PILImage from app.services import storage from app.services.llm.client import LLMResult def _png_bytes(size=(40, 40)) -> bytes: buf = io.BytesIO() PILImage.new("RGB", size, color="white").save(buf, "PNG") return buf.getvalue() class TestImageValidation: def test_accepts_png(self): mime, ext = storage.validate_image(_png_bytes()) assert (mime, ext) == ("image/png", "png") def test_rejects_non_image(self): with pytest.raises(storage.ImageValidationError): storage.validate_image(b"this is not an image") def test_rejects_oversized(self, monkeypatch): monkeypatch.setattr(storage.settings, "max_upload_bytes", 10) with pytest.raises(storage.ImageValidationError, match="大小上限"): storage.validate_image(_png_bytes()) def test_rejects_huge_dimensions(self, monkeypatch): """防解压炸弹：像素维度上限。""" monkeypatch.setattr(storage, "_MAX_DIMENSION", 20) with pytest.raises(storage.ImageValidationError, match="尺寸过大"): storage.validate_image(_png_bytes(size=(50, 50))) async def _seed_job(client, headers): """建一张图 + 一个任务，返回 (job_id, user_id, object_key)。""" from sqlalchemy import select from app.db import SessionLocal from app.models.image import Image from app.models.job import ProcessingJob from app.models.user import User, UserSettings from app.services import crypto async with SessionLocal() as db: user = (await db.scalars(select(User))).first() db.add( UserSettings( user_id=user.id, llm_base_url="http://fake", llm_api_key_encrypted=crypto.encrypt("sk-test"), llm_text_model="text-model", llm_vision_model="vision-model", ) ) image = Image( user_id=user.id, object_key="k/1.png", mime="image/png", size_bytes=100, sha256="deadbeef", ) db.add(image) await db.flush() job = ProcessingJob(user_id=user.id, image_id=image.id, kind="ocr_extract") db.add(job) await db.commit() return job.id, user.id, image.object_key async def _run_job(client, auth_headers, vision_text: str, extract_text: str): """跑一次提取流水线，返回 (job, questions)。""" from sqlalchemy import select from app.db import SessionLocal from app.models.job import ProcessingJob from app.models.question import Question from app.services.extract_pipeline import run_extract_job h = await auth_headers() job_id, user_id, key = await _seed_job(client, h) async def fake_vision(self, system, user_text, image_bytes, mime="image/jpeg", **kw): return LLMResult(vision_text, "vision-model", 10, 5) async def fake_text(self, system, user, json_mode=False, model=None): return LLMResult(extract_text, "text-model", 8, 4) with patch( "app.services.storage.download_image", return_value=_png_bytes() ), patch( "app.services.llm.client.LLMClient.chat_vision", fake_vision ), patch( "app.services.llm.client.LLMClient.chat_text", fake_text ): await run_extract_job(job_id, user_id, key, "image/png") async with SessionLocal() as db: job = await db.get(ProcessingJob, job_id) ids = job.result_question_ids or [] questions = [] if ids: rows = ( await db.scalars(select(Question).where(Question.id.in_(ids))) ).all() by_id = {q.id: q for q in rows} # 保持流水线产出的顺序 questions = [by_id[i] for i in ids if i in by_id] return job, questions class TestExtractPipeline: async def test_single_question_extraction(self, client, auth_headers): """未配置 Pix2Text 时走多模态，成功提取出结构化笔记。""" extract = json.dumps( { "questions": [ { "type": "single_choice", "stem_markdown": "求 $x^2=9$ 的正根", "options": [ {"key": "A", "text_markdown": "2"}, {"key": "B", "text_markdown": "3"}, ], "correct_answer": ["B"], "difficulty": 2, } ] } ) job, qs = await _run_job(client, auth_headers, "求 $x^2=9$ 的正根", extract) assert job.status.value == "done" assert job.engine_used == "vlm" # 降级到多模态 assert len(qs) == 1 assert qs[0].type.value == "single_choice" assert qs[0].correct_answer == ["B"] assert qs[0].ocr_engine == "vlm" async def test_one_image_splits_into_many_notes(self, client, auth_headers): """一张照片常含多道题，应拆成多条独立笔记并共享原图。""" extract = json.dumps( { "questions": [ { "type": "single_choice", "stem_markdown": "求 $x^2=9$ 的正根", "options": [{"key": "A", "text_markdown": "3"}], "correct_answer": ["A"], }, { "type": "true_false", "stem_markdown": "偶函数关于 y 轴对称", "options": [ {"key": "T", "text_markdown": "正确"}, {"key": "F", "text_markdown": "错误"}, ], "correct_answer": ["T"], }, { "type": "short_answer", "stem_markdown": "简述牛顿第一定律", "correct_answer": None, }, ] } ) job, qs = await _run_job(client, auth_headers, "三道题的 OCR 文本", extract) assert job.status.value == "done" assert len(job.result_question_ids) == 3 assert len(qs) == 3 assert [q.type.value for q in qs] == [ "single_choice", "true_false", "short_answer", ] # 都指向同一张原图 image_ids = {q.source_image_id for q in qs} assert len(image_ids) == 1 and None not in image_ids async def test_original_ocr_text_is_cached(self, client, auth_headers): """原始 OCR 文本要留在 image 上，便于日后追溯/重新拆分。""" from app.db import SessionLocal from app.models.image import Image as ImageModel extract = json.dumps( {"questions": [{"type": "short_answer", "stem_markdown": "第一题"}]} ) job, _ = await _run_job(client, auth_headers, "完整的 OCR 原文", extract) async with SessionLocal() as db: image = await db.get(ImageModel, job.image_id) assert image.ocr_markdown == "完整的 OCR 原文" async def test_malformed_json_falls_back_to_unclassified( self, client, auth_headers ): """整体无法解析时不丢内容：存成一条未分类笔记，原文进题干。""" job, qs = await _run_job( client, auth_headers, "一些无法结构化的题目文字", "这不是 JSON" ) assert job.status.value == "done" assert len(qs) == 1 assert qs[0].type.value == "unclassified" assert qs[0].stem_markdown == "一些无法结构化的题目文字" async def test_empty_questions_array_falls_back(self, client, auth_headers): job, qs = await _run_job( client, auth_headers, "原始文本", json.dumps({"questions": []}) ) assert len(qs) == 1 assert qs[0].type.value == "unclassified" assert qs[0].stem_markdown == "原始文本" async def test_one_bad_item_does_not_lose_the_others(self, client, auth_headers): """某一条解析失败，不该拖垮同一张图里的其他题。""" extract = json.dumps( { "questions": [ { "type": "single_choice", "stem_markdown": "好的题", "options": [{"key": "A", "text_markdown": "1"}], "correct_answer": ["A"], }, {"type": "不存在的题型", "stem_markdown": "题型认不出"}, ] } ) job, qs = await _run_job(client, auth_headers, "OCR 原文", extract) assert job.status.value == "done" assert len(qs) == 2 assert qs[0].type.value == "single_choice" # 认不出的题型退化为未分类，但内容保住了 assert qs[1].type.value == "unclassified" assert qs[1].stem_markdown == "题型认不出" async def test_json_in_code_fence_is_parsed(self, client, auth_headers): """模型常把 JSON 包在 ```json 代码块里，要能容忍。""" fenced = ( '```json\n{"questions": [{"type": "short_answer", ' '"stem_markdown": "简述牛顿第一定律"}]}\n```' ) job, qs = await _run_job(client, auth_headers, "简述牛顿第一定律", fenced) assert job.status.value == "done" assert len(qs) == 1 assert qs[0].stem_markdown == "简述牛顿第一定律" async def test_extracted_notes_need_no_answer(self, client, auth_headers): """图上没答案是常态，不该因此失败或硬填。""" extract = json.dumps( { "questions": [ { "type": "short_answer", "stem_markdown": "证明勾股定理", "correct_answer": None, } ] } ) _, qs = await _run_job(client, auth_headers, "证明勾股定理", extract) assert qs[0].correct_answer is None async def test_missing_ai_config_fails_job_gracefully(self, client, auth_headers): """用户没配 AI 时任务应标记失败并给出可读原因，而不是卡住。""" from sqlalchemy import select from app.db import SessionLocal from app.models.image import Image from app.models.job import ProcessingJob from app.models.user import User from app.services.extract_pipeline import run_extract_job await auth_headers() async with SessionLocal() as db: user = (await db.scalars(select(User))).first() image = Image( user_id=user.id, object_key="k/2.png", mime="image/png", size_bytes=1, sha256="beef", ) db.add(image) await db.flush() job = ProcessingJob(user_id=user.id, image_id=image.id) db.add(job) await db.commit() job_id, user_id = job.id, user.id # 没有 UserSettings → LLMConfigError await run_extract_job(job_id, user_id, "k/2.png", "image/png") async with SessionLocal() as db: job = await db.get(ProcessingJob, job_id) assert job.status.value == "failed" assert "AI" in job.error_message class TestOcrFallback: async def test_pix2text_unhealthy_falls_back_to_vlm(self, monkeypatch): """配了 Pix2Text 但服务不健康时，应自动降级到多模态。""" from app.services.ocr.pipeline import run_ocr monkeypatch.setattr( "app.services.ocr.pipeline.settings.ocr_service_url", "http://ocr:8001" ) class FakeLLM: async def chat_vision(self, *a, **kw): return LLMResult("vlm output", "vision-model", 1, 1) async def unhealthy(self): return False with patch( "app.services.ocr.pix2text_client.Pix2TextClient.healthy", unhealthy ): result = await run_ocr(b"fake", "image/png", FakeLLM()) assert result.engine == "vlm" assert result.markdown == "vlm output" async def test_pix2text_used_when_healthy(self, monkeypatch): from app.services.ocr.base import OcrResult from app.services.ocr.pipeline import run_ocr monkeypatch.setattr( "app.services.ocr.pipeline.settings.ocr_service_url", "http://ocr:8001" ) async def healthy(self): return True async def recognize(self, image_bytes, mime): return OcrResult(markdown="pix2text output", confidence=0.9, engine="pix2text") with patch( "app.services.ocr.pix2text_client.Pix2TextClient.healthy", healthy ), patch( "app.services.ocr.pix2text_client.Pix2TextClient.recognize", recognize ): result = await run_ocr(b"fake", "image/png", None) assert result.engine == "pix2text" async def test_low_confidence_falls_back_to_vlm(self, monkeypatch): """Pix2Text 置信度过低时也要降级。""" from app.services.ocr.base import OcrResult from app.services.ocr.pipeline import run_ocr monkeypatch.setattr( "app.services.ocr.pipeline.settings.ocr_service_url", "http://ocr:8001" ) class FakeLLM: async def chat_vision(self, *a, **kw): return LLMResult("vlm output", "vision-model", 1, 1) async def healthy(self): return True async def low_conf(self, image_bytes, mime): return OcrResult(markdown="garbled", confidence=0.1, engine="pix2text") with patch( "app.services.ocr.pix2text_client.Pix2TextClient.healthy", healthy ), patch( "app.services.ocr.pix2text_client.Pix2TextClient.recognize", low_conf ): result = await run_ocr(b"fake", "image/png", FakeLLM()) assert result.engine == "vlm"