Dish_Recognition/tools/test_media_upload.py

# coding=utf-8
"""本地对象存储 + 上传接口 + 先传图再上报 冒烟。"""
import json
import os
import sys
import threading
import time
from pathlib import Path
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

TMP = ROOT / "_tmp_media_test"


def _jpeg(path):
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.new("RGB", (8, 8), (10, 20, 30)).save(path, "JPEG")
    return path


def test_store():
    from central_server.media_store import MediaStore

    store = MediaStore(root=str(TMP / "store"), public_base="http://127.0.0.1:9000/media")
    src = _jpeg(TMP / "in.jpg")
    saved = store.save_file(str(src), media_type="source", base_id="BASE_T", request_id="REQ_T")
    assert saved and os.path.isfile(saved["path"])
    assert "/media/" in saved["url"]
    print("store ok")


def test_server():
    os.environ["DISH_MEDIA_ROOT"] = str(TMP / "srv")
    os.environ["DISH_MEDIA_PUBLIC_BASE"] = "http://127.0.0.1:19090/media"
    os.environ["DISH_RESULT_STORE"] = str(TMP / "results")
    os.environ.pop("DISH_API_TOKEN", None)

    import central_server.api_server as api
    from central_server.media_store import MediaStore

    api.RESULT_DIR = os.environ["DISH_RESULT_STORE"]
    os.makedirs(api.RESULT_DIR, exist_ok=True)
    api._store = MediaStore(
        root=os.environ["DISH_MEDIA_ROOT"],
        public_base=os.environ["DISH_MEDIA_PUBLIC_BASE"],
    )

    threading.Thread(target=lambda: api.serve("127.0.0.1", 19090), daemon=True).start()
    time.sleep(0.5)

    health = json.loads(urlopen("http://127.0.0.1:19090/api/health", timeout=3).read())
    assert health["code"] == 0
    print("health ok")

    src = _jpeg(TMP / "cap.jpg")
    boundary = "----TestBoundary"
    chunks = []
    def field(name, value):
        chunks.append(f"--{boundary}\r\nContent-Disposition: form-data; name=\"{name}\"\r\n\r\n{value}\r\n".encode())
    field("type", "source")
    field("base_id", "BASE_T")
    field("request_id", "REQ_HTTP")
    chunks.append(
        f"--{boundary}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"cap.jpg\"\r\n"
        "Content-Type: image/jpeg\r\n\r\n".encode()
    )
    chunks.append(src.read_bytes())
    chunks.append(f"\r\n--{boundary}--\r\n".encode())
    body = b"".join(chunks)
    req = Request(
        "http://127.0.0.1:19090/api/dish/media/upload",
        data=body,
        method="POST",
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
    )
    up = json.loads(urlopen(req, timeout=5).read())
    assert up["code"] == 0 and up["data"]["url"]
    print("upload ok", up["data"]["url"])
    img_bytes = urlopen(up["data"]["url"], timeout=3).read()
    assert len(img_bytes) > 10
    print("static ok")

    payload = {
        "request_id": "REQ_HTTP",
        "base_id": "BASE_T",
        "capture_time": "2026-09-15 14:00:00",
        "source_image": up["data"]["url"],
        "labeled_image": None,
        "dish_count": 1,
        "dishes": [{
            "index": 1, "name": None, "score": 0.1, "is_unknown": True,
            "suggested_name": "x", "box": None, "crop_image": None, "text": "未知",
        }],
    }
    req2 = Request(
        "http://127.0.0.1:19090/api/dish/recognize/result",
        data=json.dumps(payload).encode(),
        method="POST",
        headers={"Content-Type": "application/json"},
    )
    res = json.loads(urlopen(req2, timeout=5).read())
    assert res["code"] == 0 and res["data"]["accepted"] and not res["data"]["duplicate"]
    res2 = json.loads(urlopen(req2, timeout=5).read())
    assert res2["data"]["duplicate"] is True
    bad = dict(payload)
    bad["dish_count"] = 9
    req3 = Request(
        "http://127.0.0.1:19090/api/dish/recognize/result",
        data=json.dumps(bad).encode(),
        method="POST",
        headers={"Content-Type": "application/json"},
    )
    try:
        urlopen(req3, timeout=5)
        raise AssertionError("expected 400")
    except HTTPError as exc:
        assert exc.code == 400
    print("result ok")


def test_client_urlize():
    from central_server.media_store import MediaStore
    from service.result_api import ResultApiClient, build_result_payload

    src = _jpeg(TMP / "client.jpg")
    client = ResultApiClient(mock=True, base_id="BASE_T")
    client._store = MediaStore(
        root=str(TMP / "client_media"),
        public_base="http://127.0.0.1:19090/media",
    )
    payload = build_result_payload(
        str(src),
        [{"is_unknown": False, "label": "土豆丝", "score": 0.9, "crop_image": str(src)}],
        labeled_image=str(src),
        base_id="BASE_T",
        request_id="REQ_CLIENT",
    )
    out = client._ensure_remote_urls(payload)
    assert out["source_image"].startswith("http")
    assert out["dishes"][0]["crop_image"].startswith("http")
    print("client urlize ok")


if __name__ == "__main__":
    test_store()
    test_server()
    test_client_urlize()
    print("ALL OK")