pyr0ball
a6f60c9e07
Full FastAPI backend for the AI music continuation editor:
Services
- chain.py: chain + node CRUD, commit/discard, recursive CTE spine query
- musicgen.py: MusicGenClient with cf-orch allocation + mock mode (CF_MUSICGEN_MOCK=1)
- stems.py: Demucs 4-stem separation subprocess wrapper + mock mode
- export.py: ffmpeg concat demuxer to stitch committed spine into WAV/MP3
API endpoints
- chains: CRUD, multipart audio upload (WAV/MP3/FLAC/OGG/M4A/AIFF)
- nodes: branch creation (202 + BackgroundTasks), commit, discard, audio stream
- gpu: cf-orch capacity status; session allocation stubbed pending cf-orch#43
- stems: Paid-tier stem separation (Demucs, gated via tiers.py)
- export: POST /{chain_id}/export → FileResponse download
- events: SSE stream (node-status events) per chain via asyncio Queue pub/sub
Infrastructure
- lifespan: reads SPARROW_DB_PATH/DATA_DIR at startup (not import time)
- events_store: subscribe/unsubscribe/broadcast pattern for SSE
- CORS: open in dev, SPARROW_CORS_ORIGINS in production
- Background generation opens its own DB connection (WAL-safe)
Tests: 30/30 passing across service units and API integration
150 lines
4.9 KiB
Python
150 lines
4.9 KiB
Python
# tests/test_chain_service.py — unit tests for chain/node CRUD service
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from app.services import chain as svc
|
|
|
|
|
|
def test_create_and_list_chain(conn):
|
|
chain = svc.create_chain(conn, "test chain")
|
|
assert chain["name"] == "test chain"
|
|
assert chain["node_count"] == 0
|
|
|
|
chains = svc.list_chains(conn)
|
|
assert len(chains) == 1
|
|
assert chains[0]["id"] == chain["id"]
|
|
|
|
|
|
def test_get_chain_not_found(conn):
|
|
assert svc.get_chain(conn, "nonexistent") is None
|
|
|
|
|
|
def test_delete_chain(conn):
|
|
chain = svc.create_chain(conn, "deleteme")
|
|
assert svc.delete_chain(conn, chain["id"])
|
|
assert svc.get_chain(conn, chain["id"]) is None
|
|
|
|
|
|
def test_create_root_node(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "root test")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close() # empty file, just needs to exist
|
|
|
|
node = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
assert node["status"] == "ready"
|
|
assert node["is_committed"] is True
|
|
assert node["parent_id"] is None
|
|
assert node["audio_path"] == audio
|
|
assert node["duration_s"] == 30.0
|
|
|
|
|
|
def test_create_branch_node(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "branch test")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close()
|
|
root = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
|
|
branch = svc.create_branch_node(
|
|
conn,
|
|
parent_id=root["id"],
|
|
chain_id=chain["id"],
|
|
prompt="upbeat jazz",
|
|
energy=0.8,
|
|
tempo_feel=None,
|
|
density=None,
|
|
cfg_coef=3.0,
|
|
prompt_duration_s=10.0,
|
|
)
|
|
assert branch["status"] == "pending"
|
|
assert branch["is_committed"] is False
|
|
assert branch["parent_id"] == root["id"]
|
|
assert branch["prompt"] == "upbeat jazz"
|
|
|
|
|
|
def test_update_node_status(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "status test")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close()
|
|
root = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
branch = svc.create_branch_node(
|
|
conn, root["id"], chain["id"], "test", None, None, None, 3.0, 10.0
|
|
)
|
|
|
|
updated = svc.update_node_status(
|
|
conn, branch["id"], "generating"
|
|
)
|
|
assert updated["status"] == "generating"
|
|
|
|
out_path = str(tmp_path / "out.wav")
|
|
open(out_path, "wb").close()
|
|
done = svc.update_node_status(
|
|
conn, branch["id"], "ready", audio_path=out_path, duration_s=15.0
|
|
)
|
|
assert done["status"] == "ready"
|
|
assert done["audio_path"] == out_path
|
|
assert done["duration_s"] == 15.0
|
|
|
|
|
|
def test_commit_node_discards_siblings(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "commit test")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close()
|
|
root = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
|
|
b1 = svc.create_branch_node(
|
|
conn, root["id"], chain["id"], "branch 1", None, None, None, 3.0, 10.0
|
|
)
|
|
b2 = svc.create_branch_node(
|
|
conn, root["id"], chain["id"], "branch 2", None, None, None, 3.0, 10.0
|
|
)
|
|
|
|
# Make both ready
|
|
for nid in (b1["id"], b2["id"]):
|
|
svc.update_node_status(conn, nid, "ready",
|
|
audio_path=audio, duration_s=15.0)
|
|
|
|
# Commit b1 — b2 should be deleted
|
|
committed = svc.commit_node(conn, b1["id"])
|
|
assert committed["is_committed"] is True
|
|
assert svc.get_node(conn, b2["id"]) is None
|
|
|
|
|
|
def test_delete_node_refuses_root(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "del test")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close()
|
|
root = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
assert svc.delete_node(conn, root["id"]) is False
|
|
|
|
|
|
def test_delete_node_refuses_committed(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "del committed")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close()
|
|
root = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
branch = svc.create_branch_node(
|
|
conn, root["id"], chain["id"], "", None, None, None, 3.0, 10.0
|
|
)
|
|
svc.update_node_status(conn, branch["id"], "ready",
|
|
audio_path=audio, duration_s=15.0)
|
|
svc.commit_node(conn, branch["id"])
|
|
assert svc.delete_node(conn, branch["id"]) is False
|
|
|
|
|
|
def test_committed_spine_order(conn, tmp_path):
|
|
chain = svc.create_chain(conn, "spine test")
|
|
audio = str(tmp_path / "audio.wav")
|
|
open(audio, "wb").close()
|
|
root = svc.create_root_node(conn, chain["id"], audio, 30.0)
|
|
|
|
child = svc.create_branch_node(
|
|
conn, root["id"], chain["id"], "child", None, None, None, 3.0, 10.0
|
|
)
|
|
svc.update_node_status(conn, child["id"], "ready",
|
|
audio_path=audio, duration_s=15.0)
|
|
svc.commit_node(conn, child["id"])
|
|
|
|
spine = svc.get_committed_spine(conn, chain["id"])
|
|
assert len(spine) == 2
|
|
assert spine[0]["id"] == root["id"]
|
|
assert spine[1]["id"] == child["id"]
|