初始化提交

tianmu/tests/test_health.py

@@ -0,0 +1,78 @@
+"""
+天目健康检查测试（pytest）。
+
+需先启动天目服务：
+  cd tianmu && source .venv/bin/activate && python3 -m src.main
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import subprocess
+import time
+import urllib.request
+
+import pytest
+
+API_BASE = os.environ.get("TIANMU_API_BASE", "http://localhost:8101")
+
+
+def test_healthz():
+    """健康检查应返回 200 并包含引擎状态。"""
+    req = urllib.request.Request(f"{API_BASE}/healthz")
+    with urllib.request.urlopen(req, timeout=5) as resp:
+        assert resp.status == 200
+        data = json.loads(resp.read())
+        assert "status" in data
+        assert "engines" in data
+        assert "version" in data
+        assert data["version"] == "0.2.0"
+        assert "ppocrv6" in data["engines"]
+
+
+def test_ocr_empty_image():
+    """发送空数据应得到有效的错误响应（而非崩溃）。"""
+    body = json.dumps({"image": ""}).encode("utf-8")
+    req = urllib.request.Request(
+        f"{API_BASE}/v1/ocr",
+        data=body,
+        headers={"Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(req, timeout=5) as resp:
+        assert resp.status == 200
+        data = json.loads(resp.read())
+        # 空图片不会触发 422，而是返回 engine 处理结果或错误
+        assert "engine" in data
+        assert "error" in data  # 可能是 None 或错误消息，但字段必须存在
+
+
+def test_ocr_png_image():
+    """PP-OCRv6 应正确识别简单图片中的文字。"""
+    from PIL import Image, ImageDraw
+
+    # 生成白底黑字图片
+    img = Image.new("RGB", (400, 60), color="white")
+    draw = ImageDraw.Draw(img)
+    draw.text((10, 20), "Hello World 123", fill="black")
+
+    import io
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+
+    body = json.dumps({"image": b64, "engine": "ppocrv6"}).encode("utf-8")
+    req = urllib.request.Request(
+        f"{API_BASE}/v1/ocr",
+        data=body,
+        headers={"Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(req, timeout=30) as resp:
+        data = json.loads(resp.read())
+        assert data["engine"] == "ppocrv6"
+        assert data["error"] is None
+        assert len(data["text"]) > 0  # 识别到文字
+        assert len(data["blocks"]) > 0
+        assert data["blocks"][0]["confidence"] > 0.5
+        assert data["duration_ms"] > 0