79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
"""
|
|
天目健康检查测试（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
|