tests/test_main.py

import io
import uuid

import pytest

pytestmark = pytest.mark.asyncio


async def test_health(client):
    response = await client.get("/health")
    assert response.status_code == 200
    assert response.json() == {"status": "ok"}


async def test_status_unknown_job(client):
    response = await client.get("/status/nonexistent-job-id")
    assert response.status_code == 404


async def test_status_known_job(client, monkeypatch):
    import processor
    from models import JobState, JobStatus

    fake_status = JobStatus(job_id="abc", state=JobState.pending, progress=0)
    monkeypatch.setattr(processor, "_jobs", {"abc": fake_status})

    response = await client.get("/status/abc")
    assert response.status_code == 200
    data = response.json()
    assert data["job_id"] == "abc"
    assert data["state"] == "pending"
    assert data["progress"] == 0
    assert data["error"] is None


async def test_upload_mp3(client):
    fake_mp3 = io.BytesIO(b"\xff\xfb" + b"\x00" * 100)
    response = await client.post(
        "/upload",
        files={"file": ("song.mp3", fake_mp3, "audio/mpeg")},
    )
    assert response.status_code == 200
    data = response.json()
    assert "job_id" in data
    uuid.UUID(data["job_id"])


async def test_upload_with_trim_params(client):
    fake_mp3 = io.BytesIO(b"\xff\xfb" + b"\x00" * 100)
    response = await client.post(
        "/upload",
        files={"file": ("song.mp3", fake_mp3, "audio/mpeg")},
        data={"trim_start_sec": "4.5", "trim_end_sec": "22.75"},
    )
    assert response.status_code == 200


async def test_upload_rejects_invalid_trim_range(client):
    fake_mp3 = io.BytesIO(b"\xff\xfb" + b"\x00" * 100)
    response = await client.post(
        "/upload",
        files={"file": ("song.mp3", fake_mp3, "audio/mpeg")},
        data={"trim_start_sec": "12", "trim_end_sec": "10"},
    )
    assert response.status_code == 400
    assert "greater" in response.json()["detail"].lower()


async def test_upload_rejects_negative_trim_start(client):
    fake_mp3 = io.BytesIO(b"\xff\xfb" + b"\x00" * 100)
    response = await client.post(
        "/upload",
        files={"file": ("song.mp3", fake_mp3, "audio/mpeg")},
        data={"trim_start_sec": "-1"},
    )
    assert response.status_code == 400


async def test_upload_rejects_non_mp3(client):
    fake_file = io.BytesIO(b"fake content")
    response = await client.post(
        "/upload",
        files={"file": ("doc.pdf", fake_file, "application/pdf")},
    )
    assert response.status_code == 400
    assert "MP3" in response.json()["detail"]


async def test_upload_accepts_audio_mp3_content_type(client):
    fake_mp3 = io.BytesIO(b"\xff\xfb" + b"\x00" * 100)
    response = await client.post(
        "/upload",
        files={"file": ("song.mp3", fake_mp3, "audio/mp3")},
    )
    assert response.status_code == 200


async def test_upload_rejects_oversized_file(client, monkeypatch):
    import main

    monkeypatch.setattr(main, "MAX_FILE_SIZE", 8)
    oversized = io.BytesIO(b"\xff\xfb" + b"\x00" * 20)
    response = await client.post(
        "/upload",
        files={"file": ("big.mp3", oversized, "audio/mpeg")},
    )
    assert response.status_code == 400
    assert "large" in response.json()["detail"].lower()


async def test_upload_rejects_non_mp3_bytes(client):
    fake_zip = io.BytesIO(b"PK\x03\x04" + b"\x00" * 100)
    response = await client.post(
        "/upload",
        files={"file": ("trap.mp3", fake_zip, "audio/mpeg")},
    )
    assert response.status_code == 400
    assert "valid MP3" in response.json()["detail"]


async def test_download_done_job(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "done-job-123"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    fake_output = job_dir / "instrumental.mp3"
    fake_output.write_bytes(b"\xff\xfb" + b"\x00" * 200)

    fake_status = JobStatus(
        job_id=job_id,
        state=JobState.done,
        progress=100,
        stems={"instrumental": "instrumental.mp3"},
    )
    monkeypatch.setitem(processor._jobs, job_id, fake_status)
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)
    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    monkeypatch.setattr(main, "DOWNLOAD_LOG_FILE", tmp_path / "download_log.json")

    response = await client.get(f"/download/{job_id}")
    assert response.status_code == 200
    assert response.headers["content-type"].startswith("audio/mpeg")
    assert response.content == b"\xff\xfb" + b"\x00" * 200


async def test_download_specific_stem(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "stem-job"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    vocal_file = job_dir / "vocals.mp3"
    vocal_file.write_bytes(b"\xff\xfb" + b"\x00" * 60)

    fake_status = JobStatus(
        job_id=job_id,
        state=JobState.done,
        progress=100,
        stems={"vocals": "vocals.mp3"},
    )
    monkeypatch.setitem(processor._jobs, job_id, fake_status)
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)
    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    monkeypatch.setattr(main, "DOWNLOAD_LOG_FILE", tmp_path / "download_log.json")

    response = await client.get(f"/download/{job_id}/vocals")
    assert response.status_code == 200
    assert response.content == b"\xff\xfb" + b"\x00" * 60




async def test_download_original_stem(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "original-job"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    original_file = job_dir / "original.mp3"
    original_file.write_bytes(b"\xff\xfb" + b"\x00" * 70)

    fake_status = JobStatus(
        job_id=job_id,
        state=JobState.done,
        progress=100,
        stems={"instrumental": "instrumental.mp3"},
    )
    monkeypatch.setitem(processor._jobs, job_id, fake_status)
    monkeypatch.setattr(main, "UPLOAD_DIR", tmp_path / "uploads")
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)

    response = await client.get(f"/download/{job_id}/original")
    assert response.status_code == 200
    assert response.content == b"\xff\xfb" + b"\x00" * 70


async def test_download_original_stem_falls_back_to_upload_dir(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "original-fallback-job"
    upload_dir = tmp_path / "uploads"
    upload_dir.mkdir(parents=True, exist_ok=True)
    upload_file = upload_dir / f"{job_id}.mp3"
    upload_file.write_bytes(b"\xff\xfb" + b"\x00" * 70)

    fake_status = JobStatus(
        job_id=job_id,
        state=JobState.done,
        progress=100,
        stems={"instrumental": "instrumental.mp3"},
    )
    monkeypatch.setitem(processor._jobs, job_id, fake_status)
    monkeypatch.setattr(main, "UPLOAD_DIR", upload_dir)
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)

    response = await client.get(f"/download/{job_id}/original")
    assert response.status_code == 200
    assert response.content == b"\xff\xfb" + b"\x00" * 70


async def test_download_unknown_stem_returns_404(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "unknown-stem-job"
    fake_status = JobStatus(
        job_id=job_id,
        state=JobState.done,
        progress=100,
        stems={"vocals": "vocals.mp3"},
    )
    monkeypatch.setitem(processor._jobs, job_id, fake_status)
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)

    response = await client.get(f"/download/{job_id}/drums")
    assert response.status_code == 404


async def test_download_processing_job_returns_409(client, monkeypatch):
    import processor
    from models import JobState, JobStatus

    job_id = "processing-job"
    fake_status = JobStatus(job_id=job_id, state=JobState.processing, progress=40)
    monkeypatch.setitem(processor._jobs, job_id, fake_status)

    response = await client.get(f"/download/{job_id}")
    assert response.status_code == 409
    assert "complete" in response.json()["detail"].lower()


async def test_download_unknown_job_returns_404(client):
    response = await client.get("/download/does-not-exist")
    assert response.status_code == 404


async def test_get_stats_returns_zeros_when_no_file(client, tmp_path, monkeypatch):
    import main

    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    response = await client.get("/stats")
    assert response.status_code == 200
    data = response.json()
    assert data["total_visits"] == 100
    assert data["total_downloads"] == 5


async def test_record_visit_increments_total_visits(client, tmp_path, monkeypatch):
    import main

    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    response = await client.post("/stats/visit")
    assert response.status_code == 204
    stats = (await client.get("/stats")).json()
    assert stats["total_visits"] == 101


async def test_record_visit_is_cumulative(client, tmp_path, monkeypatch):
    import main

    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    await client.post("/stats/visit")
    await client.post("/stats/visit")
    stats = (await client.get("/stats")).json()
    assert stats["total_visits"] == 102


async def test_download_rejects_out_of_range_pitch(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "pitch-range-job"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    (job_dir / "instrumental.mp3").write_bytes(b"\xff\xfb" + b"\x00" * 100)

    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)
    monkeypatch.setitem(
        processor._jobs,
        job_id,
        JobStatus(
            job_id=job_id,
            state=JobState.done,
            progress=100,
            stems={"instrumental": "instrumental.mp3"},
        ),
    )

    response = await client.get(f"/download/{job_id}?pitch_semitones=20")
    assert response.status_code == 400
    assert "pitch" in response.json()["detail"].lower()


async def test_download_rejects_out_of_range_tempo(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "tempo-range-job"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    (job_dir / "instrumental.mp3").write_bytes(b"\xff\xfb" + b"\x00" * 100)

    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)
    monkeypatch.setitem(
        processor._jobs,
        job_id,
        JobStatus(
            job_id=job_id,
            state=JobState.done,
            progress=100,
            stems={"instrumental": "instrumental.mp3"},
        ),
    )

    response = await client.get(f"/download/{job_id}?tempo_pct=25")
    assert response.status_code == 400
    assert "tempo" in response.json()["detail"].lower()


async def test_download_default_params_serve_raw_file(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "raw-serve-job"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    raw_bytes = b"\xff\xfb" + b"\x00" * 123
    (job_dir / "instrumental.mp3").write_bytes(raw_bytes)

    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    monkeypatch.setattr(main, "DOWNLOAD_LOG_FILE", tmp_path / "download_log.json")
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)
    monkeypatch.setitem(
        processor._jobs,
        job_id,
        JobStatus(
            job_id=job_id,
            state=JobState.done,
            progress=100,
            stems={"instrumental": "instrumental.mp3"},
        ),
    )

    response = await client.get(
        f"/download/{job_id}?pitch_semitones=0&tempo_pct=100"
    )
    assert response.status_code == 200
    assert response.content == raw_bytes


async def test_download_increments_total_downloads(client, tmp_path, monkeypatch):
    import main
    import processor
    from models import JobState, JobStatus

    job_id = "dl-stat-job"
    job_dir = tmp_path / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    fake_output = job_dir / "instrumental.mp3"
    fake_output.write_bytes(b"\xff\xfb" + b"\x00" * 100)

    monkeypatch.setattr(main, "STATS_FILE", tmp_path / "stats.json")
    monkeypatch.setattr(main, "DOWNLOAD_LOG_FILE", tmp_path / "download_log.json")
    monkeypatch.setattr(main, "OUTPUT_DIR", tmp_path)
    monkeypatch.setitem(
        processor._jobs,
        job_id,
        JobStatus(
            job_id=job_id,
            state=JobState.done,
            progress=100,
            stems={"instrumental": "instrumental.mp3"},
        ),
    )

    await client.get(f"/download/{job_id}")
    stats = (await client.get("/stats")).json()
    assert stats["total_downloads"] == 6