35 lines
2.0 KiB
Python
35 lines
2.0 KiB
Python
|
import time
|
||
|
|
|
||
|
|
from spriteforge.studio.jobs import GenerationQueue
|
||
|
|
from spriteforge.studio.models import Asset, GenerationRecipe, GenerationRequest
|
||
|
|
from spriteforge.studio.store import ProjectStore
|
||
|
|
|
||
|
|
|
||
|
|
def test_generation_queue_persists_status_and_candidates(tmp_path):
|
||
|
|
store=ProjectStore.create(tmp_path/"studio","ops","Ops")
|
||
|
|
store.add_asset(Asset(id="agent",name="Agent",kind="character",target_width=60,target_height=60))
|
||
|
|
_,shot=store.add_shot("agent","idle",0,0)
|
||
|
|
request=GenerationRequest(project_id="ops",asset_id="agent",shot_id=shot.id,count=2,
|
||
|
|
recipe=GenerationRecipe(provider="diagnostic",model="test",seed=1,prompt="x",width=64,height=64))
|
||
|
|
queue=GenerationQueue(store);job=queue.submit(request);deadline=time.monotonic()+3
|
||
|
|
while time.monotonic()<deadline:
|
||
|
|
saved=next(x for x in store.load().jobs if x.id==job.id)
|
||
|
|
if saved.status not in {"queued","running"}:break
|
||
|
|
time.sleep(.01)
|
||
|
|
queue.close()
|
||
|
|
assert saved.status=="succeeded" and len(saved.candidate_ids)==2
|
||
|
|
assert len(store.load().assets[0].shots[0].candidates)==2
|
||
|
|
|
||
|
|
|
||
|
|
def test_cancel_retry_and_interrupted_recovery(tmp_path):
|
||
|
|
store=ProjectStore.create(tmp_path/"studio","ops","Ops");store.add_asset(Asset(id="agent",name="Agent",kind="character",target_width=60,target_height=60));_,shot=store.add_shot("agent","idle",0,0)
|
||
|
|
request=GenerationRequest(project_id="ops",asset_id="agent",shot_id=shot.id,count=1,recipe=GenerationRecipe(provider="missing",model="x",seed=1,prompt="x",width=64,height=64))
|
||
|
|
job=store.create_job(request);store.update_job(job.id,status="running")
|
||
|
|
queue=GenerationQueue(store)
|
||
|
|
assert store.get_job(job.id).status=="failed" and "stopped" in store.get_job(job.id).error
|
||
|
|
retry=queue.retry(job.id);deadline=time.monotonic()+2
|
||
|
|
while time.monotonic()<deadline and store.get_job(retry.id).status in {"queued","running"}:time.sleep(.01)
|
||
|
|
assert store.get_job(retry.id).status=="failed"
|
||
|
|
queued=store.create_job(request);queue.cancel(queued.id);assert store.get_job(queued.id).status=="cancelled"
|
||
|
|
queue.close()
|