89 lines
5.8 KiB
Python
89 lines
5.8 KiB
Python
import concurrent.futures,io,json,sqlite3,threading,time,uuid,hashlib
|
|
from email import policy
|
|
from email.parser import BytesParser
|
|
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
|
|
from pathlib import Path
|
|
import requests
|
|
from PIL import Image
|
|
|
|
data=Path('D:/web/dj/integration-work/hd-test-references');base='http://127.0.0.1:18886/api/v1'
|
|
calls=[]
|
|
class Mock(BaseHTTPRequestHandler):
|
|
def log_message(self,*args):pass
|
|
def do_POST(self):
|
|
raw=self.rfile.read(int(self.headers['Content-Length']));content_type=self.headers.get('Content-Type','')
|
|
fields={};images=[]
|
|
if 'multipart' in content_type:
|
|
msg=BytesParser(policy=policy.default).parsebytes(('Content-Type: '+content_type+'\r\n\r\n').encode()+raw)
|
|
for part in msg.iter_parts():
|
|
name=part.get_param('name',header='content-disposition');value=part.get_payload(decode=True)
|
|
if name=='image[]':
|
|
with Image.open(io.BytesIO(value)) as im:images.append(list(im.convert('RGB').getpixel((50,50))))
|
|
else:fields[name]=value.decode()
|
|
else:fields=json.loads(raw)
|
|
calls.append({'path':self.path,'fields':fields,'images':images})
|
|
response=json.dumps({'data':[{'url':'http://127.0.0.1:18887/result.png'}]}).encode()
|
|
self.send_response(200);self.send_header('Content-Type','application/json');self.end_headers();self.wfile.write(response)
|
|
def do_GET(self):
|
|
raw=(data/'green.png').read_bytes();self.send_response(200);self.send_header('Content-Type','image/png');self.end_headers();self.wfile.write(raw)
|
|
server=ThreadingHTTPServer(('127.0.0.1',18887),Mock)
|
|
threading.Thread(target=server.serve_forever,daemon=True).start()
|
|
def sql(query,args=()):
|
|
with sqlite3.connect(data/'huobao.sqlite3') as db:
|
|
db.row_factory=sqlite3.Row;return [dict(x) for x in db.execute(query,args)]
|
|
def upload(kind,id,color='red'):
|
|
with (data/(color+'.png')).open('rb') as f:r=requests.post(f'{base}/{kind}/{id}/reference-images',files={'file':(color+'.png',f,'image/png')},timeout=10)
|
|
return r
|
|
def refs(kind,id):return json.loads(sql(f'SELECT reference_images FROM {kind} WHERE id=?',(id,))[0]['reference_images'] or '[]')
|
|
def remove(kind,id,path):
|
|
r=requests.delete(f'{base}/{kind}/{id}/reference-images',json={'path':path});assert r.status_code==200,r.text
|
|
def generate(kind,id,cfg=14):
|
|
r=requests.post(f'{base}/{kind}/{id}/generate-image',json={'episode_id':7,'model':'qwen_image_2_1','config_id':cfg},timeout=20)
|
|
assert r.status_code==200,r.text
|
|
tid=r.json()['data']['image_generation_id'];return wait(tid)
|
|
def wait(tid):
|
|
for _ in range(100):
|
|
row=sql('SELECT * FROM sys_task WHERE id=?',(tid,))[0]
|
|
if row['status']!='processing':return row
|
|
time.sleep(.1)
|
|
raise AssertionError(('task timeout',tid))
|
|
try:
|
|
before={table:sql(f'SELECT image_url,local_path FROM {table} WHERE id=?',(id,))[0] for table,id in [('characters',26),('scenes',22),('props',1)]}
|
|
for table,id in [('characters',26),('scenes',22),('props',1)]:
|
|
for color in ['red','blue']:
|
|
r=upload(table,id,color);assert r.status_code==200,r.text
|
|
assert len(refs(table,id))==2
|
|
assert sql(f'SELECT image_url,local_path FROM {table} WHERE id=?',(id,))[0]==before[table]
|
|
row=generate(table,id);assert row['status']=='completed',row['error_msg']
|
|
assert len(json.loads(row['params'])['referenceImages'])==2
|
|
call=calls[-1];assert call['path']=='/v1/images/edits' and call['fields']['gpu']=='gpu1',call
|
|
assert len(call['images'])==2 and call['images'][0][0]>180 and call['images'][1][2]>180,call
|
|
print('PASS: character/scene/prop persist references, preserve current output on upload, send ordered multipart images with GPU1')
|
|
# Independent assets and per-character batch snapshots.
|
|
assert upload('characters',27,'green').status_code==200
|
|
r=requests.post(base+'/characters/batch-generate-images',json={'character_ids':[26,27],'episode_id':7,'model':'qwen_image_2_1','config_id':16},timeout=20)
|
|
assert r.status_code==200,r.text
|
|
for tid in r.json()['data']['ids']:assert wait(tid)['status']=='completed'
|
|
batch=calls[-2:];assert sorted(len(c['images']) for c in batch)==[1,2] and all(c['fields']['gpu']=='gpu3' for c in batch)
|
|
print('PASS: batch retains different per-character references and selected GPU3')
|
|
# Removing references does not erase files already snapshotted into past tasks.
|
|
saved=refs('characters',26);oldpaths=sql('SELECT image_url,local_path FROM characters WHERE id=26')[0]
|
|
for path in saved:remove('characters',26,path);assert (data/path).is_file()
|
|
assert refs('characters',26)==[] and sql('SELECT image_url,local_path FROM characters WHERE id=26')[0]==oldpaths
|
|
assert generate('characters',26)['status']=='completed' and calls[-1]['path']=='/v1/images/generations'
|
|
# Invalid uploads, unknown/deleted assets and concurrent cap enforcement.
|
|
r=requests.post(base+'/characters/26/reference-images',files={'file':('fake.png',b'not an image','image/png')});assert r.status_code==400
|
|
assert upload('characters',999999).status_code==404
|
|
for _ in range(5):assert upload('characters',26).status_code==200
|
|
with concurrent.futures.ThreadPoolExecutor() as pool:statuses=list(pool.map(lambda _:upload('characters',26).status_code,range(2)))
|
|
assert sorted(statuses)==[200,400] and len(refs('characters',26))==6,statuses
|
|
# One unavailable reference must fail visibly, never fall back to text-only.
|
|
broken=refs('characters',26)[0];(data/broken).rename(data/'temporarily-missing.png');count=len(calls)
|
|
row=generate('characters',26);assert row['status']=='failed' and '参考图读取失败' in row['error_msg']
|
|
assert len(calls)==count
|
|
(data/'temporarily-missing.png').rename(data/broken)
|
|
for path in refs('characters',26)[2:]:remove('characters',26,path)
|
|
print('PASS: removal preserves output/history files, empty refs use text generation, invalid files rejected, concurrent cap enforced, missing refs never silently ignored')
|
|
(data/'integration-result.json').write_text(json.dumps({'passed':True,'calls':len(calls),'asset_types':3,'maximum':6},indent=2))
|
|
finally:server.shutdown()
|