chengma/cmshoppe
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(product-suite): add direct job state machine

Browse Source
This commit is contained in:
chengma
2026-07-20 18:12:44 +08:00
parent 2a2cbad7bd
commit aa000ff5e6
11 changed files with 820 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_image_studio_generation.py
+234 -22
View File
@@ -4,6 +4,7 @@ import sys
import threading
import time
import unittest
import base64
from concurrent.futures import CancelledError
from unittest import mock
@@ -11,7 +12,7 @@ sys.path.insert(0, os.path.dirname(__file__))
from _helpers import TempDirMixin
from app import db, image_studio, image_studio_generation
from app import ai, appconfig, db, image_studio, image_studio_generation
class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
@@ -74,6 +75,51 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
"download_with_curl": "false",
}
def _direct_config(self, temp_dir):
cfg = self._config(temp_dir)
models_path = os.path.join(temp_dir, "ai_models.json")
cfg["ai"]["backend"] = "direct"
cfg["ai"]["default_text_model"] = "Text"
cfg["ai"]["default_image_model"] = "Direct Image"
cfg["ai_models_path"] = models_path
appconfig.save_ai_models_config(
{
"models": [
{
"name": "Text",
"category": "text",
"enabled": True,
"url": "https://text.example.com/v1",
"model": "text-model",
"api_key": "sk-text",
"api_type": "chat",
"connect_timeout_seconds": 3,
"timeout_seconds": 10,
"extra_body": {},
},
{
"name": "Direct Image",
"category": "image",
"enabled": True,
"url": "https://image.example.com/v1",
"model": "image-model",
"api_key": "sk-image",
"api_type": "images_edits",
"connect_timeout_seconds": 7,
"timeout_seconds": 12,
"extra_body": {},
},
]
},
path=models_path,
)
return ai.freeze_runtime_config(
cfg,
models_path=models_path,
include_cmhub=False,
include_direct_models=True,
)
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)
@@ -149,27 +195,28 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_direct_gateway_rejects_new_suite_job_before_creation(self):
def test_direct_gateway_creates_direct_suite_job_before_execution(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
cfg["ai"]["backend"] = "direct"
with mock.patch("app.image_studio_generation.image_studio.create_job") as create_job:
with self.assertRaises(image_studio_generation.ImageStudioGenerationError):
image_studio_generation.create_generation_jobs(
project.id,
source.id,
"不应提交",
1,
config=cfg,
path=cfg["db_path"],
)
jobs = image_studio_generation.create_generation_jobs(
project.id,
source.id,
"提交前来源快照",
1,
config=cfg,
path=cfg["db_path"],
)
create_job.assert_not_called()
self.assertEqual(1, len(jobs))
self.assertEqual("direct", jobs[0].generation_source)
self.assertEqual("openai_images_edits", jobs[0].provider)
self.assertIsNone(jobs[0].task_id)
self.assert_removed(temp_dir)
def test_resume_rejects_non_default_gateway_task_without_reading_gateway_config(self):
def test_direct_job_with_invalid_task_id_is_not_sent_to_any_gateway(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
cfg["ai"]["backend"] = "direct"
@@ -179,16 +226,22 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
job_type="白底图",
prompt="旧任务",
generation_source="direct",
provider="direct",
path=cfg["db_path"],
)
job = image_studio.set_job_submitted(
job.id,
"custom-task-1",
provider="openai_images_edits",
path=cfg["db_path"],
)
conn = db.connect(cfg["db_path"])
try:
conn.execute(
"UPDATE image_studio_jobs SET task_id = 'custom-task-1' WHERE id = ?",
(job.id,),
)
conn.commit()
finally:
conn.close()
job = image_studio.get_job(job.id, path=cfg["db_path"])
with mock.patch("app.image_studio_generation._runtime") as runtime:
with mock.patch("app.image_studio_generation._runtime") as runtime, \
mock.patch("app.image_studio_generation._direct_runtime") as direct_runtime:
summary = image_studio_generation.run_jobs(
[job],
config=cfg,
@@ -196,9 +249,168 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
)
runtime.assert_not_called()
direct_runtime.assert_not_called()
self.assertEqual(1, summary["total"])
self.assertEqual(1, summary["failed"])
self.assertIn("不属于默认网关", summary["jobs"][0]["error"])
self.assertIn("来源无效", summary["jobs"][0]["error"])
self.assert_removed(temp_dir)
def test_direct_job_uses_one_ordered_edit_request_and_records_output_dimensions(self):
with self.make_temp_dir() as temp_dir:
_cfg, project, source = self._project_source(temp_dir)
cfg = self._direct_config(temp_dir)
reference_path = os.path.join(temp_dir, "reference.png")
with open(reference_path, "wb") as fh:
fh.write(self._png_bytes())
reference = image_studio.add_asset(
project.id,
image_studio.ASSET_KIND_ORIGINAL,
local_path=reference_path,
source_order=2,
path=cfg["db_path"],
)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
reference_asset_ids=[reference.id],
job_type="场景图",
prompt="已冻结的套图提示词",
generation_source="direct",
provider="openai_images_edits",
path=cfg["db_path"],
)
seen = []
encoded = base64.b64encode(self._png_bytes()).decode("ascii")
def fake_direct_call(model, body, config, attempts, **kwargs):
seen.append((model, body, config, attempts, kwargs))
return {"data": [{"b64_json": encoded}]}
with mock.patch(
"app.image_studio_generation.ai._call_with_retry",
side_effect=fake_direct_call,
), mock.patch("app.image_studio_generation._runtime") as cmhub_runtime, \
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as cmhub_submit, \
mock.patch("app.image_studio_generation.ai._cmhub_call_once") as cmhub_poll:
summary = image_studio_generation.run_jobs(
[job],
aspect_ratio="3:4",
config=cfg,
path=cfg["db_path"],
run_session_id="direct-test-session",
)
self.assertEqual(1, summary["success"])
self.assertEqual("1024x1536", summary["output"]["requested_output_size"])
self.assertTrue(summary["output"]["approximate_ratio"])
self.assertEqual(1, len(seen))
_model, body, _config, attempts, kwargs = seen[0]
self.assertEqual(1, attempts)
self.assertEqual("multipart", kwargs["request_kind"])
self.assertEqual(2, body.count(b'name="image[]"'))
self.assertLess(body.index(b"source.png"), body.index(b"reference.png"))
self.assertIn(b'name="n"', body)
self.assertIn(b"\r\n1\r\n", body)
cmhub_runtime.assert_not_called()
cmhub_submit.assert_not_called()
cmhub_poll.assert_not_called()
stored = image_studio.get_job(job.id, path=cfg["db_path"])
self.assertEqual("succeeded", stored.status)
self.assertEqual("direct", stored.generation_source)
self.assertEqual("openai_images_edits", stored.provider)
self.assertIsNone(stored.task_id)
self.assertEqual("direct-test-session", stored.run_session_id)
asset = image_studio.get_asset(stored.output_asset_id, path=cfg["db_path"])
self.assertIsNone(asset.remote_url)
self.assertEqual("3:4", asset.aspect_ratio)
self.assertEqual("1024x1536", asset.requested_output_size)
self.assertEqual((1024, 1536), (asset.rendered_width, asset.rendered_height))
self.assert_removed(temp_dir)
def test_direct_job_saves_returned_image_after_stop_and_never_retries(self):
with self.make_temp_dir() as temp_dir:
_cfg, project, source = self._project_source(temp_dir)
cfg = self._direct_config(temp_dir)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
prompt="停止后保存",
generation_source="direct",
provider="openai_images_edits",
path=cfg["db_path"],
)
stopped = {"value": False}
encoded = base64.b64encode(self._png_bytes()).decode("ascii")
def fake_direct_call(*args, **kwargs):
self.assertEqual(1, args[3])
stopped["value"] = True
return {"data": [{"b64_json": encoded}]}
with mock.patch(
"app.image_studio_generation.ai._call_with_retry",
side_effect=fake_direct_call,
) as direct_call:
summary = image_studio_generation.run_jobs(
[job],
config=cfg,
path=cfg["db_path"],
should_stop=lambda: stopped["value"],
)
self.assertEqual(1, summary["success"])
self.assertEqual(1, direct_call.call_count)
self.assertEqual("succeeded", image_studio.get_job(job.id, path=cfg["db_path"]).status)
self.assert_removed(temp_dir)
def test_startup_recovery_marks_only_stale_direct_running_jobs_failed(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
stale = image_studio.create_job(
project.id,
source_asset_id=source.id,
prompt="已中断",
generation_source="direct",
provider="openai_images_edits",
path=cfg["db_path"],
)
active = image_studio.create_job(
project.id,
source_asset_id=source.id,
prompt="仍在执行",
generation_source="direct",
provider="openai_images_edits",
path=cfg["db_path"],
)
image_studio.update_job_status(
stale.id,
"running",
run_session_id="previous-session",
path=cfg["db_path"],
)
image_studio.update_job_status(
active.id,
"running",
run_session_id="active-session",
path=cfg["db_path"],
)
recovered = image_studio.fail_stale_direct_jobs(
active_run_session_ids=["active-session"],
path=cfg["db_path"],
)
self.assertEqual([stale.id], [job.id for job in recovered])
stale = image_studio.get_job(stale.id, path=cfg["db_path"])
active = image_studio.get_job(active.id, path=cfg["db_path"])
self.assertEqual("failed", stale.status)
self.assertIn("程序中断", stale.error)
self.assertEqual(image_studio.JOB_RECOVERY_REGENERATE, stale.recovery_action)
self.assertEqual("running", active.status)
self.assert_removed(temp_dir)