30 lines
1.1 KiB
Python
30 lines
1.1 KiB
Python
import os, tempfile, importlib
|
|
|
|
def _fresh_db(monkeypatch, tmp):
|
|
monkeypatch.setenv("IMAGE_DATA_DIR", tmp)
|
|
import app.db as db
|
|
importlib.reload(db)
|
|
db.init_db()
|
|
return db
|
|
|
|
def test_create_and_get_task(monkeypatch):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
db = _fresh_db(monkeypatch, tmp)
|
|
row = db.create_task("t1", "gpt_image", {"prompt": "a cat"})
|
|
assert row["id"] == "t1"
|
|
assert row["provider"] == "gpt_image"
|
|
assert row["status"] == "queued"
|
|
got = db.get_task("t1")
|
|
assert got["id"] == "t1"
|
|
assert db.get_task("nope") is None
|
|
|
|
def test_update_task_sets_image_url(monkeypatch):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
db = _fresh_db(monkeypatch, tmp)
|
|
db.create_task("t2", "nano_banana", {"prompt": "x"})
|
|
db.update_task("t2", "succeeded", 100, message="done", image_url="/media/image/t2.png")
|
|
got = db.get_task("t2")
|
|
assert got["status"] == "succeeded"
|
|
assert got["image_url"] == "/media/image/t2.png"
|
|
assert got["progress"] == 100
|