drum-sample-extractor / scripts /test_selected_export_card_actions.py
ChatGPT
feat: add spleeter and selected card exports
fa35534
Raw
History Blame Contribute Delete
3.36 kB
#!/usr/bin/env python3
"""Smoke-test card selection semantics, selected-only export, draw-next, and immediate clip edits."""
from __future__ import annotations
import io
import json
import sys
import time
import zipfile
from pathlib import Path
from urllib.parse import quote
import soundfile as sf
from fastapi.testclient import TestClient
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from app import app # noqa: E402
from synth_generator import generate_test_song # noqa: E402
def wait_for_job(client: TestClient, job_id: str) -> dict:
for _ in range(120):
payload = client.get(f"/api/jobs/{job_id}").json()
if payload["status"] in {"complete", "error"}:
return payload
time.sleep(0.1)
raise TimeoutError(job_id)
def main() -> int:
song = generate_test_song(pattern_name="funk", bars=1, bpm=124, add_bass=False)
buf = io.BytesIO()
sf.write(buf, song.drums_only, song.sr, format="WAV")
buf.seek(0)
client = TestClient(app)
response = client.post(
"/api/jobs",
files={"file": ("cards.wav", buf, "audio/wav")},
data={"params": json.dumps({"stem": "all", "clustering_mode": "online_preview", "target_min": 3, "target_max": 10})},
)
response.raise_for_status()
job_id = response.json()["id"]
job = wait_for_job(client, job_id)
assert job["status"] == "complete", job.get("error")
samples = job["result"]["samples"]
assert samples, "expected at least one sample"
labels = [sample["label"] for sample in samples[: min(2, len(samples))]]
selected = client.post(f"/api/jobs/{job_id}/export-selected", json={"labels": labels})
selected.raise_for_status()
selected_payload = selected.json()["export"]
assert selected_payload["kind"] == "selected-sample-export"
assert selected_payload["selected_labels"] == sorted(labels)
archive_response = client.get(selected_payload["file_urls"]["archive"])
archive_response.raise_for_status()
with zipfile.ZipFile(io.BytesIO(archive_response.content)) as zf:
names = zf.namelist()
assert any(name.endswith(".wav") for name in names), names
label = labels[0]
draw = client.post(f"/api/jobs/{job_id}/samples/{quote(label, safe='')}/draw", json={})
draw.raise_for_status()
drawn = draw.json()["sample"]
assert drawn["label"] == label
assert drawn["url"]
drawn_audio = client.get(drawn["url"])
drawn_audio.raise_for_status()
assert drawn_audio.content[:4] == b"RIFF"
edit = client.post(f"/api/jobs/{job_id}/samples/{quote(label, safe='')}/edit", json={"start_offset_ms": 5, "tail_offset_ms": 30})
edit.raise_for_status()
edited = edit.json()["sample"]
assert edited["label"] == label
assert "overrides/hits" in edited["file"], edited
edited_audio = client.get(edited["url"])
edited_audio.raise_for_status()
assert edited_audio.content[:4] == b"RIFF"
print(json.dumps({
"status": "ok",
"job_id": job_id,
"selected_labels": labels,
"drawn_representative_hit_index": drawn["representative_hit_index"],
"edited_file": edited["file"],
"archive": selected_payload["files"]["archive"],
}, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())