451 lines
19 KiB
Python
451 lines
19 KiB
Python
import io
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
from _helpers import TempDirMixin
|
|
|
|
from app import db, image_studio, image_studio_generation
|
|
|
|
|
|
class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
|
|
def _png_bytes(self):
|
|
from PIL import Image
|
|
|
|
output = io.BytesIO()
|
|
Image.new("RGB", (32, 32), (120, 80, 160)).save(output, format="PNG")
|
|
return output.getvalue()
|
|
|
|
def _config(self, temp_dir):
|
|
return {
|
|
"data_dir": temp_dir,
|
|
"db_path": os.path.join(temp_dir, "cmshopee.db"),
|
|
"image_dir": os.path.join(temp_dir, "images"),
|
|
"ai": {
|
|
"backend": "cmhub",
|
|
"image_concurrency": 1,
|
|
"retry": 0,
|
|
"resolution": "1k",
|
|
"jpg_quality": 90,
|
|
"cmhub": {
|
|
"base_url": "https://cmhub.example.com",
|
|
"image_alias": "image-hd",
|
|
"connect_timeout": 3,
|
|
"download_with_curl": "false",
|
|
},
|
|
},
|
|
}
|
|
|
|
def _project_source(self, temp_dir):
|
|
cfg = self._config(temp_dir)
|
|
db.init_db(cfg["db_path"])
|
|
project = image_studio.create_or_get_project(
|
|
account_alias="alias",
|
|
account_slug="alias_slug",
|
|
item_id="51100639510",
|
|
path=cfg["db_path"],
|
|
)
|
|
source_path = os.path.join(temp_dir, "source.png")
|
|
with open(source_path, "wb") as fh:
|
|
fh.write(self._png_bytes())
|
|
source = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
remote_url="https://cdn.example.com/source.png",
|
|
local_path=source_path,
|
|
source_order=1,
|
|
path=cfg["db_path"],
|
|
)
|
|
return cfg, project, source
|
|
|
|
def _runtime(self):
|
|
return {
|
|
"base_url": "https://cmhub.example.com",
|
|
"api_key": "sk-test",
|
|
"alias": "image-hd",
|
|
"connect_timeout": 3,
|
|
"use_system_proxy": False,
|
|
"download_with_curl": "false",
|
|
}
|
|
|
|
def test_generate_image_jobs_creates_independent_jobs_and_assets(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
submitted = []
|
|
poll_seen_persisted = []
|
|
|
|
def fake_submit(method, url, api_key, **kwargs):
|
|
self.assertEqual("POST", method)
|
|
task_id = "cmhub-task-%d" % (len(submitted) + 1)
|
|
submitted.append((task_id, kwargs["headers_extra"]["Idempotency-Key"]))
|
|
return {
|
|
"task_id": task_id,
|
|
"status": "queued",
|
|
"call_id": "call-%d" % len(submitted),
|
|
"points_cost": 2,
|
|
"points_balance": 100 - len(submitted) * 2,
|
|
}
|
|
|
|
def fake_poll(method, url, api_key, **kwargs):
|
|
self.assertEqual("GET", method)
|
|
task_id = url.rsplit("/", 1)[-1]
|
|
jobs = image_studio.list_resumable_jobs(path=cfg["db_path"], project_id=project.id)
|
|
poll_seen_persisted.append(any(job.task_id == task_id for job in jobs))
|
|
return {
|
|
"task_id": task_id,
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/%s.png" % task_id},
|
|
}
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
|
side_effect=fake_submit,
|
|
), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_once",
|
|
side_effect=fake_poll,
|
|
), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
|
return_value=(self._png_bytes(), 0.1),
|
|
), \
|
|
mock.patch("app.image_studio_generation.ai._sleep_cmhub_poll"):
|
|
summary = image_studio_generation.generate_image_jobs(
|
|
project.id,
|
|
source.id,
|
|
"完整提示词",
|
|
3,
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(3, summary["success"])
|
|
self.assertEqual(0, summary["failed"])
|
|
self.assertEqual(3, len(submitted))
|
|
self.assertEqual(3, len({key for _, key in submitted}))
|
|
self.assertEqual([True, True, True], poll_seen_persisted)
|
|
jobs = image_studio.list_resumable_jobs(path=cfg["db_path"], project_id=project.id)
|
|
self.assertEqual([], jobs)
|
|
all_jobs = [
|
|
image_studio.get_job(result["job"].id, path=cfg["db_path"])
|
|
for result in summary["jobs"]
|
|
]
|
|
self.assertEqual(["succeeded", "succeeded", "succeeded"], [job.status for job in all_jobs])
|
|
self.assertTrue(all(job.task_id for job in all_jobs))
|
|
self.assertTrue(all(job.task_key for job in all_jobs))
|
|
self.assertEqual([2, 2, 2], [job.points_cost for job in all_jobs])
|
|
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
|
|
self.assertEqual(3, len(assets))
|
|
self.assertTrue(all(os.path.isfile(asset.local_path) for asset in assets))
|
|
self.assertTrue(all(asset.parent_asset_id == source.id for asset in assets))
|
|
self.assertTrue(all(asset.prompt == "完整提示词" for asset in assets))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_generate_image_jobs_sends_selected_aspect_ratio(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
submitted_payloads = []
|
|
|
|
def fake_submit(method, url, api_key, **kwargs):
|
|
submitted_payloads.append(dict(kwargs["payload"]))
|
|
return {"task_id": "cmhub-ratio", "status": "queued"}
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
|
side_effect=fake_submit,
|
|
), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_once",
|
|
return_value={
|
|
"task_id": "cmhub-ratio",
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/ratio.png"},
|
|
},
|
|
), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
|
return_value=(self._png_bytes(), 0.1),
|
|
):
|
|
summary = image_studio_generation.generate_image_jobs(
|
|
project.id,
|
|
source.id,
|
|
"比例测试",
|
|
1,
|
|
aspect_ratio="3:4",
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(1, summary["success"])
|
|
self.assertEqual("3:4", submitted_payloads[0]["aspect_ratio"])
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_stop_after_download_discards_temporary_result(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
stopped = {"value": False}
|
|
saved_paths = []
|
|
|
|
def fake_download(request_result, out_path, config, on_event, job_id):
|
|
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
|
with open(out_path, "wb") as fh:
|
|
fh.write(self._png_bytes())
|
|
saved_paths.append(out_path)
|
|
stopped["value"] = True
|
|
return out_path
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
|
return_value={"task_id": "cmhub-stop", "status": "queued"},
|
|
), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_once",
|
|
return_value={
|
|
"task_id": "cmhub-stop",
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/stop.png"},
|
|
},
|
|
), \
|
|
mock.patch(
|
|
"app.image_studio_generation._download_and_save_job_image",
|
|
side_effect=fake_download,
|
|
):
|
|
summary = image_studio_generation.generate_image_jobs(
|
|
project.id,
|
|
source.id,
|
|
"停止测试",
|
|
1,
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
should_stop=lambda: stopped["value"],
|
|
)
|
|
|
|
self.assertEqual(1, summary["cancelled"])
|
|
self.assertEqual([], image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"]))
|
|
self.assertEqual(1, len(saved_paths))
|
|
self.assertFalse(os.path.exists(saved_paths[0]))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_resume_existing_job_polls_without_new_submit(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
job_type="main",
|
|
prompt="续查提示词",
|
|
path=cfg["db_path"],
|
|
)
|
|
image_studio.set_job_submitted(job.id, "cmhub-task-resume", path=cfg["db_path"])
|
|
|
|
def fake_poll(method, url, api_key, **kwargs):
|
|
return {
|
|
"task_id": "cmhub-task-resume",
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/resume.png"},
|
|
}
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit, \
|
|
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
|
return_value=(self._png_bytes(), 0.1),
|
|
):
|
|
summary = image_studio_generation.resume_image_jobs(
|
|
project_id=project.id,
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(1, summary["success"])
|
|
submit.assert_not_called()
|
|
updated = image_studio.get_job(job.id, path=cfg["db_path"])
|
|
self.assertEqual("succeeded", updated.status)
|
|
self.assertEqual("cmhub-task-resume", updated.task_id)
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_download_failure_does_not_submit_again_or_create_asset(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
submit_count = 0
|
|
|
|
def fake_submit(method, url, api_key, **kwargs):
|
|
nonlocal submit_count
|
|
submit_count += 1
|
|
return {"task_id": "cmhub-task-1", "status": "queued"}
|
|
|
|
def fake_poll(method, url, api_key, **kwargs):
|
|
return {
|
|
"task_id": "cmhub-task-1",
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/fail.png"},
|
|
}
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
|
side_effect=fake_submit,
|
|
), \
|
|
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
|
side_effect=RuntimeError("download failed"),
|
|
):
|
|
summary = image_studio_generation.generate_image_jobs(
|
|
project.id,
|
|
source.id,
|
|
"完整提示词",
|
|
1,
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(0, summary["success"])
|
|
self.assertEqual(1, summary["failed"])
|
|
self.assertEqual(1, submit_count)
|
|
jobs = [result["job"] for result in summary["jobs"]]
|
|
self.assertEqual("failed", image_studio.get_job(jobs[0].id, path=cfg["db_path"]).status)
|
|
self.assertEqual([], image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"]))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_resume_failed_download_job_polls_without_new_submit(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
job_type="main",
|
|
prompt="恢复下载",
|
|
path=cfg["db_path"],
|
|
)
|
|
image_studio.set_job_submitted(
|
|
job.id,
|
|
"cmhub-task-download",
|
|
call_id="call-download",
|
|
points_cost=2,
|
|
points_balance=88,
|
|
path=cfg["db_path"],
|
|
)
|
|
image_studio.update_job_status(
|
|
job.id,
|
|
"failed",
|
|
error="下载新封面失败",
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
failed_job = image_studio.get_job(job.id, path=cfg["db_path"])
|
|
self.assertEqual(image_studio.JOB_RECOVERY_RESUME, failed_job.recovery_action)
|
|
|
|
resumable = image_studio.list_resumable_jobs(
|
|
path=cfg["db_path"],
|
|
project_id=project.id,
|
|
include_failed_downloads=True,
|
|
)
|
|
self.assertEqual([job.id], [item.id for item in resumable])
|
|
|
|
def fake_poll(method, url, api_key, **kwargs):
|
|
self.assertEqual("GET", method)
|
|
return {
|
|
"task_id": "cmhub-task-download",
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/recovered.png"},
|
|
"points_balance": 88,
|
|
}
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit, \
|
|
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
|
return_value=(self._png_bytes(), 0.1),
|
|
):
|
|
summary = image_studio_generation.resume_image_jobs(
|
|
project_id=project.id,
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(1, summary["success"])
|
|
submit.assert_not_called()
|
|
updated = image_studio.get_job(job.id, path=cfg["db_path"])
|
|
self.assertEqual("succeeded", updated.status)
|
|
self.assertEqual("cmhub-task-download", updated.task_id)
|
|
self.assertEqual("call-download", updated.call_id)
|
|
self.assertEqual(image_studio.JOB_RECOVERY_NONE, updated.recovery_action)
|
|
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
|
|
self.assertEqual(1, len(assets))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_failed_cmhub_task_marks_only_that_job_failed(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, source = self._project_source(temp_dir)
|
|
submitted = []
|
|
|
|
def fake_submit(method, url, api_key, **kwargs):
|
|
task_id = "cmhub-task-%d" % (len(submitted) + 1)
|
|
submitted.append(task_id)
|
|
return {"task_id": task_id, "status": "queued"}
|
|
|
|
def fake_poll(method, url, api_key, **kwargs):
|
|
task_id = url.rsplit("/", 1)[-1]
|
|
if task_id.endswith("-2"):
|
|
return {"task_id": task_id, "status": "failed", "error": {"message": "上游失败"}}
|
|
return {
|
|
"task_id": task_id,
|
|
"status": "succeeded",
|
|
"result": {"image_url": "https://cdn.example.com/%s.png" % task_id},
|
|
}
|
|
|
|
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
|
side_effect=fake_submit,
|
|
), \
|
|
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
|
mock.patch(
|
|
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
|
return_value=(self._png_bytes(), 0.1),
|
|
):
|
|
summary = image_studio_generation.generate_image_jobs(
|
|
project.id,
|
|
source.id,
|
|
"完整提示词",
|
|
2,
|
|
config=cfg,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(1, summary["success"])
|
|
self.assertEqual(1, summary["failed"])
|
|
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
|
|
self.assertEqual(1, len(assets))
|
|
failed_jobs = [
|
|
job
|
|
for job in image_studio.list_resumable_jobs(
|
|
path=cfg["db_path"],
|
|
project_id=project.id,
|
|
include_failed_downloads=True,
|
|
)
|
|
if job.status == "failed"
|
|
]
|
|
self.assertEqual([], failed_jobs)
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|