File size: 4,507 Bytes
f7d8394 6f7704a f7d8394 6f7704a f7d8394 6f7704a f7d8394 6f7704a f7d8394 6f7704a f7d8394 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | import io
import zipfile
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from PIL import Image
def make_png_bytes(color=(255, 0, 0, 255), size=(8, 8)):
img = Image.new("RGBA", size, color)
buf = io.BytesIO()
img.save(buf, format="PNG")
return buf.getvalue()
class FakeInf:
def __init__(self):
self.killed = 0
async def remove_bg(self, data, label=None):
return make_png_bytes(color=(0, 255, 0, 128))
def queue_depth(self):
return 0
def in_flight(self):
return 0
def jobs(self):
return [
{"id": "abc", "status": "processing", "label": "a.png",
"age_seconds": 1.2},
]
def kill_all(self):
self.killed += 1
return {"killed_pending": 0, "worker_restarted": True}
@pytest.fixture
def client(monkeypatch):
import api
fake = FakeInf()
monkeypatch.setattr(api, "inf", fake)
app = FastAPI()
app.include_router(api.router)
return TestClient(app)
def test_health(client):
r = client.get("/health")
assert r.status_code == 200
body = r.json()
assert body["status"] == "ok"
assert body["model"] == "birefnet-general"
assert body["queue"] == 0
assert body["in_flight"] == 0
def test_jobs(client):
r = client.get("/jobs")
assert r.status_code == 200
body = r.json()
assert body["total"] == 1
assert "jobs" in body
assert body["jobs"][0]["status"] == "processing"
assert body["jobs"][0]["label"] == "a.png"
def test_remove_single(client):
r = client.post(
"/remove",
files={"file": ("a.png", make_png_bytes(), "image/png")},
)
assert r.status_code == 200
assert r.headers["content-type"] == "image/png"
assert r.content[:8] == b"\x89PNG\r\n\x1a\n"
def test_remove_rejects_non_image(client):
r = client.post(
"/remove",
files={"file": ("a.txt", b"hello", "text/plain")},
)
assert r.status_code == 400
def test_remove_batch_returns_zip(client):
files = [
("files", ("a.png", make_png_bytes(), "image/png")),
("files", ("b.png", make_png_bytes(), "image/png")),
]
r = client.post("/remove/batch", files=files)
assert r.status_code == 200
assert r.headers["content-type"] == "application/zip"
zf = zipfile.ZipFile(io.BytesIO(r.content))
names = zf.namelist()
assert "a.png" in names
assert "b.png" in names
def test_remove_batch_rejects_too_many(client, monkeypatch):
import config
monkeypatch.setattr(config, "MAX_BATCH_FILES", 1)
files = [
("files", ("a.png", make_png_bytes(), "image/png")),
("files", ("b.png", make_png_bytes(), "image/png")),
]
r = client.post("/remove/batch", files=files)
assert r.status_code == 400
def test_remove_batch_records_per_file_error(client, monkeypatch):
import api
calls = {"n": 0}
orig = api.inf.remove_bg
async def flaky(data, label=None):
calls["n"] += 1
if calls["n"] == 1:
raise RuntimeError("boom")
return await orig(data, label=label)
monkeypatch.setattr(api.inf, "remove_bg", flaky)
files = [
("files", ("bad.png", make_png_bytes(), "image/png")),
("files", ("good.png", make_png_bytes(), "image/png")),
]
r = client.post("/remove/batch", files=files)
assert r.status_code == 200
zf = zipfile.ZipFile(io.BytesIO(r.content))
assert "_errors.txt" in zf.namelist()
assert "good.png" in zf.namelist()
assert b"bad.png" in zf.read("_errors.txt")
def test_kill_without_configured_key_returns_503(client, monkeypatch):
import config
monkeypatch.setattr(config, "KILL_API_KEY", None)
r = client.get("/kill")
assert r.status_code == 503
def test_kill_wrong_key_returns_401(client, monkeypatch):
import config
monkeypatch.setattr(config, "KILL_API_KEY", "secret")
r = client.get("/kill", params={"key": "nope"})
assert r.status_code == 401
def test_kill_missing_key_returns_401(client, monkeypatch):
import config
monkeypatch.setattr(config, "KILL_API_KEY", "secret")
r = client.get("/kill")
assert r.status_code == 401
def test_kill_correct_key_succeeds(client, monkeypatch):
import config
monkeypatch.setattr(config, "KILL_API_KEY", "secret")
r = client.get("/kill", params={"key": "secret"})
assert r.status_code == 200
body = r.json()
assert body["worker_restarted"] is True
|