drum-sample-extractor / scripts /test_api_job.py
ChatGPT
feat: add full-context reproduction and clearer controls
0b5f0f0
Raw
History Blame Contribute Delete
1.22 kB
import io, json, sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import soundfile as sf
from fastapi.testclient import TestClient
from app import app
from synth_generator import generate_test_song
song=generate_test_song(pattern_name='rock', bars=1, bpm=120, add_bass=False)
buf=io.BytesIO()
sf.write(buf, song.drums_only, song.sr, format='WAV')
buf.seek(0)
client=TestClient(app)
params={'stem':'all','target_min':2,'target_max':6,'synthesize':True}
r=client.post('/api/jobs', files={'file':('test.wav', buf, 'audio/wav')}, data={'params':json.dumps(params)})
r.raise_for_status()
job=r.json()
for _ in range(60):
job=client.get(f"/api/jobs/{job['id']}").json()
if job['status'] in {'complete','error'}:
break
time.sleep(0.25)
print(json.dumps({'status':job['status'], 'error':job.get('error'), 'hit_count': job.get('result',{}).get('hit_count'), 'files': job.get('result',{}).get('file_urls')}, indent=2))
assert job['status']=='complete', job.get('error')
assert job['result']['hit_count'] > 0
for key in ['source', 'stem', 'context_bed', 'target_reconstruction', 'reconstruction', 'midi', 'archive']:
assert key in job['result']['file_urls'], key