path: root/lib/project/tests/test_ai.py

"""Offline AI contract tests; no Torch installation, model download, or GPU needed."""
import hashlib
import io
import json
import os
from pathlib import Path
import sys
import time
from types import SimpleNamespace
from unittest.mock import Mock
import zipfile

import numpy as np
import pytest
import soundfile as sf

from voiceforge import ai, ai_worker, setup


MODEL_BYTES = 7986207
MODEL_HASH = "49c52edc8947ae1f9bf50d81530beaf3a2c3245aeaf34b6f31ff535cd22284d2"


@pytest.mark.parametrize("size,message", [
    (10, "checksum mismatch"),
    (MODEL_BYTES, "checksum mismatch"),
    (MODEL_BYTES + 1, "exceeds the expected size"),
])
def test_model_rejects_truncated_corrupt_and_oversized_downloads(tmp_path, monkeypatch, size, message):
    monkeypatch.setattr(ai_worker.urllib.request, "urlopen", Mock(return_value=io.BytesIO(b"x" * size)))
    progress = Mock()
    monkeypatch.setattr(ai_worker, "report", progress)
    with pytest.raises(RuntimeError, match=message):
        ai_worker.ensure_model(tmp_path / "model")
    assert list(tmp_path.iterdir()) == []
    events = [call.args for call in progress.call_args_list]
    assert events[0] == ("Downloading DeepFilterNet3 model (bytes)", 0, MODEL_BYTES)
    counts = [completed for _, completed, total in events if total == MODEL_BYTES]
    assert counts == sorted(counts)
    assert all(0 <= count <= MODEL_BYTES for count in counts)


@pytest.mark.parametrize("unsafe", [False, True])
def test_mocked_model_archive_progress_extraction_and_reuse(tmp_path, monkeypatch, unsafe):
    # Mock only the release digest for a synthetic archive; corruption tests above
    # exercise the real SHA-256 gate. Keep its real byte count and ZIP extraction.
    def archive(padding):
        stream = io.BytesIO()
        with zipfile.ZipFile(stream, "w") as bundle:
            bundle.writestr("DeepFilterNet3/config.ini", "[df]\nsr=48000\n")
            bundle.writestr("DeepFilterNet3/checkpoints/model_120.ckpt.best", "checkpoint")
            bundle.writestr("../escape" if unsafe else "padding", b"x" * padding)
        return stream.getvalue()

    payload = archive(MODEL_BYTES - len(archive(0)))
    assert len(payload) == MODEL_BYTES
    digest = Mock(wraps=hashlib.sha256())
    digest.hexdigest.return_value = MODEL_HASH
    monkeypatch.setattr(ai_worker.hashlib, "sha256", lambda: digest)
    download = Mock(return_value=io.BytesIO(payload))
    monkeypatch.setattr(ai_worker.urllib.request, "urlopen", download)
    progress = Mock()
    monkeypatch.setattr(ai_worker, "report", progress)
    directory = tmp_path / "model"
    if unsafe:
        with pytest.raises(RuntimeError, match="Unsafe path"):
            ai_worker.ensure_model(directory)
        assert not (tmp_path / "escape").exists()
        assert list(tmp_path.iterdir()) == []
    else:
        model = ai_worker.ensure_model(directory)
        assert (model / "config.ini").is_file()
        assert (model / "checkpoints/model_120.ckpt.best").is_file()
        assert ai_worker.ensure_model(directory) == model
        download.assert_called_once()
    assert b"".join(call.args[0] for call in digest.update.call_args_list) == payload
    events = [call.args for call in progress.call_args_list]
    byte_events = [event for event in events if "(bytes)" in event[0]]
    assert byte_events[0][1:] == (0, MODEL_BYTES)
    assert byte_events[-1][1:] == (MODEL_BYTES, MODEL_BYTES)
    assert len(byte_events) > 2
    assert events[-1] == ("Verifying and extracting DeepFilterNet3 model",)


def test_incomplete_model_fails_without_network(tmp_path, monkeypatch):
    directory = tmp_path / "model"
    directory.mkdir()
    download = Mock(side_effect=AssertionError("Unexpected network call"))
    monkeypatch.setattr(ai_worker.urllib.request, "urlopen", download)
    with pytest.raises(RuntimeError, match="Incomplete model directory"):
        ai_worker.ensure_model(directory)
    download.assert_not_called()


@pytest.mark.parametrize("close_output", [False, True])
def test_subprocess_heartbeat_does_not_block_on_silence_or_closed_output(close_output):
    events = []
    code = "import os,time; "
    if close_output:
        code += "os.close(1); os.close(2); "
    code += "time.sleep(0.8)"
    setup.run_process([sys.executable, "-c", code], "Waiting", lambda *event: events.append(event))
    assert events[0] == ("Waiting", None, None)
    elapsed = [completed for _, completed, total in events[1:] if total is None]
    assert len(elapsed) >= 2
    assert elapsed == sorted(elapsed)


def test_subprocess_parses_split_byte_progress_records():
    events = []
    record = "VOICEFORGE_PROGRESS " + json.dumps(["Download (bytes)", 123, 456]) + "\n"
    code = f"import os,time; os.write(1,{record[:12].encode()!r}); time.sleep(.3); os.write(1,{record[12:].encode()!r})"
    setup.run_process([sys.executable, "-c", code], "Download", lambda *event: events.append(event))
    assert events[-1] == ("Download (bytes)", 123, 456)


@pytest.mark.parametrize("exception", [KeyboardInterrupt, RuntimeError])
def test_callback_cancellation_terminates_and_reaps_child(monkeypatch, exception):
    popen = setup.subprocess.Popen
    children = []

    def launch(*args, **kwargs):
        child = popen(*args, **kwargs)
        children.append(child)
        return child

    def cancel(stage, completed, total):
        if completed is not None:
            raise exception("cancelled")

    monkeypatch.setattr(setup.subprocess, "Popen", launch)
    started = time.monotonic()
    with pytest.raises(exception, match="cancelled"):
        setup.run_process([sys.executable, "-c", "import time; time.sleep(30)"], "Waiting", cancel)
    assert len(children) == 1 and children[0].poll() is not None
    assert time.monotonic() - started < 10


def test_cancellation_kills_descendants_that_outlive_the_leader(tmp_path, monkeypatch):
    sentinel = tmp_path / "grandchild.pid"
    monkeypatch.setenv("GRANDCHILD_PID_FILE", str(sentinel))
    child_code = (
        "import os, subprocess, sys, time\n"
        "subprocess.Popen([sys.executable, '-c', "
        "\"import os, time; open(os.environ['GRANDCHILD_PID_FILE'], 'w')"
        ".write(str(os.getpid())); time.sleep(60)\"])\n"
        "while not os.path.exists(os.environ['GRANDCHILD_PID_FILE']):\n"
        "    time.sleep(0.02)\n"
    )

    def cancel(stage, completed, total):
        if completed is not None and children and children[0].poll() is not None:
            raise KeyboardInterrupt("cancelled")

    popen = setup.subprocess.Popen
    children = []

    def launch(*args, **kwargs):
        child = popen(*args, **kwargs)
        children.append(child)
        return child

    monkeypatch.setattr(setup.subprocess, "Popen", launch)
    with pytest.raises(KeyboardInterrupt, match="cancelled"):
        setup.run_process([sys.executable, "-c", child_code], "Waiting", cancel)

    def alive(pid):
        try:
            for line in Path(f"/proc/{pid}/status").read_text().splitlines():
                if line.startswith("State:"):
                    return "Z" not in line
            return True
        except FileNotFoundError:
            return False

    grandchild = int(sentinel.read_text())
    for _ in range(50):
        if not alive(grandchild):
            break
        time.sleep(0.1)
    else:
        pytest.fail("A descendant survived the cancellation of its process group")


def test_subprocess_failure_keeps_bounded_diagnostics():
    code = "import sys; print('x'*100000); print('specific failure'); sys.exit(7)"
    with pytest.raises(RuntimeError, match="exit 7") as error:
        setup.run_process([sys.executable, "-c", code], "Worker")
    assert "specific failure" in str(error.value)
    assert len(str(error.value)) < 66000


@pytest.mark.parametrize("strength", [-1, 2, float("nan"), float("inf")])
def test_invalid_strength_fails_before_setup(tmp_path, monkeypatch, strength):
    install = Mock(side_effect=AssertionError("Unexpected setup"))
    monkeypatch.setattr(ai, "ensure_ai", install)
    with pytest.raises(ValueError, match="strength"):
        ai.denoise(tmp_path / "input.wav", tmp_path / "output.wav", strength=strength)
    install.assert_not_called()


def _ready_environment(root, flavor="cpu"):
    environment = root / f"ai-df-0.5.6-torch-2.5.1-{flavor}-py311-v1"
    (environment / "bin").mkdir(parents=True)
    (environment / "bin/python").write_text("#!/bin/sh\n")
    (environment / ".voiceforge-ready").touch()
    return environment


def test_ready_environment_skips_the_check_worker(tmp_path, monkeypatch):
    root = tmp_path / "prefix"
    environment = _ready_environment(root)
    monkeypatch.setenv("VOICEFORGE_HOME", str(root))
    monkeypatch.setattr(ai.os, "confstr", lambda *_: "glibc 2.39")
    executed = []
    monkeypatch.setattr(ai, "run_process", lambda args, *rest, **kwargs: executed.append(args))
    python = ai.ensure_ai("cpu")
    assert executed == []
    assert python == environment / "bin/python"


def test_setup_verifies_even_a_ready_environment(tmp_path, monkeypatch):
    root = tmp_path / "prefix"
    _ready_environment(root)
    monkeypatch.setenv("VOICEFORGE_HOME", str(root))
    monkeypatch.setattr(ai.os, "confstr", lambda *_: "glibc 2.39")
    executed = []
    monkeypatch.setattr(ai, "run_process", lambda args, *rest, **kwargs: executed.append(args))
    ai.ensure_ai("cpu", verify=True)
    assert len(executed) == 1
    assert executed[0][-1] == "--check"


def test_broken_environment_is_rebuilt_and_self_tested(tmp_path, monkeypatch):
    root = tmp_path / "prefix"
    environment = _ready_environment(root)
    (environment / "bin/python").unlink()
    (environment / "bin/python").symlink_to("/nonexistent/python3.11")
    leftover = environment / "leftover.txt"
    leftover.write_text("from the previous location")
    monkeypatch.setenv("VOICEFORGE_HOME", str(root))
    monkeypatch.setattr(ai.os, "confstr", lambda *_: "glibc 2.39")
    stages, commands = [], []

    def fake_run(args, stage, progress=None, **kwargs):
        commands.append(args)
        stages.append(stage)
        if stage == "Checking DeepFilterNet3":
            environment.mkdir(parents=True, exist_ok=True)
            (environment / ".voiceforge-ready").touch()

    monkeypatch.setattr(ai, "run_process", fake_run)
    monkeypatch.setattr(ai, "ensure_uv", lambda progress=None: "uv")
    ai.ensure_ai("cpu")
    assert not leftover.exists()
    assert stages == ["Preparing AI Python 3.11", "Installing PyTorch (cpu)",
                      "Installing DeepFilterNet3", "Checking DeepFilterNet3"]
    assert commands[0][-1] == str(environment)
    assert (environment / ".voiceforge-ready").is_file()


def test_copied_environment_is_rebuilt_not_trusted(tmp_path, monkeypatch):
    # Copying an installation while the original remains leaves the copy's
    # interpreter links resolving into the original prefix; that must not
    # count as ready, or the copy breaks when the original is removed.
    root = tmp_path / "prefix"
    environment = _ready_environment(root)
    original = _ready_environment(tmp_path / "other")
    (environment / "bin/python").unlink()
    (environment / "bin/python").symlink_to(original / "bin/python")
    monkeypatch.setenv("VOICEFORGE_HOME", str(root))
    monkeypatch.setattr(ai.os, "confstr", lambda *_: "glibc 2.39")
    stages = []

    def fake_run(args, stage, progress=None, **kwargs):
        stages.append(stage)
        if stage == "Checking DeepFilterNet3":
            environment.mkdir(parents=True, exist_ok=True)
            (environment / ".voiceforge-ready").touch()

    monkeypatch.setattr(ai, "run_process", fake_run)
    monkeypatch.setattr(ai, "ensure_uv", lambda progress=None: "uv")
    ai.ensure_ai("cpu")
    assert stages[0] == "Preparing AI Python 3.11"


@pytest.fixture
def mocked_worker(tmp_path, monkeypatch):
    class Tensor:
        def __init__(self, array):
            self.array = array

        def unsqueeze(self, axis):
            return Tensor(np.expand_dims(self.array, axis))

        def squeeze(self, axis):
            return Tensor(np.squeeze(self.array, axis))

        def numpy(self):
            return self.array

    calls = []

    def enhance(model, state, tensor, pad):
        assert pad is True
        assert tensor.array.shape[1] % 480 == 0
        calls.append(tensor.array.shape[1])
        return Tensor(tensor.array * 0.5)

    state = SimpleNamespace(sr=lambda: 48000)
    backend = SimpleNamespace(enhance=enhance, init_df=lambda *a, **kw: (None, state, None))
    for name, module in {
        "torch": SimpleNamespace(set_num_threads=lambda n: None, from_numpy=Tensor,
                                 cuda=SimpleNamespace(is_available=lambda: False)),
        "df": SimpleNamespace(),
        "df.enhance": backend,
        "df.model": SimpleNamespace(ModelParams=lambda: SimpleNamespace(
            sr=48000, fft_size=960, hop_size=480, nb_erb=32, min_nb_freqs=2)),
        "libdf": SimpleNamespace(DF=lambda **kw: state),
    }.items():
        monkeypatch.setitem(sys.modules, name, module)
    monkeypatch.setattr(ai_worker, "ensure_model", lambda path: path)
    monkeypatch.setattr(ai_worker.signal, "signal", lambda *args: None)
    monkeypatch.setenv("DEVICE", "cpu")
    progress = Mock()
    monkeypatch.setattr(ai_worker, "report", progress)

    def run(source, target, strength):
        monkeypatch.setattr(sys, "argv", ["worker", "--device", "cpu", "--model-dir", str(tmp_path),
                                         "--source", str(source), "--target", str(target),
                                         "--strength", str(strength)])
        ai_worker.main()

    return run, calls, backend, progress


@pytest.mark.parametrize("frames", [0, 1, 479, 481, 480001, 960017])
@pytest.mark.parametrize("strength", [0, 0.85, 1])
def test_mocked_worker_exact_length_mix_and_bounded_windows(tmp_path, wav, mocked_worker, frames, strength):
    run, calls, _, progress = mocked_worker
    audio = np.random.default_rng(42).normal(0, 0.1, frames).astype("float32")
    source = wav(audio)
    target = tmp_path / "clean.wav"
    run(source, target, strength)
    result, rate = sf.read(target, dtype="float32")
    assert len(result) == frames and rate == 48000
    assert sf.info(target).subtype == "FLOAT"
    np.testing.assert_allclose(result, audio * (1 - strength * 0.5), rtol=2e-6, atol=1e-8)
    if strength == 0:
        np.testing.assert_array_equal(result, audio)
        assert calls == []
    assert max(calls, default=0) <= 674400
    assert progress.call_args.args == ("Denoising (cpu)", frames, frames)
    assert not list(tmp_path.glob(".voiceforge-ai-*"))


def test_worker_failure_preserves_existing_target(tmp_path, wav, mocked_worker):
    run, _, backend, _ = mocked_worker
    backend.enhance = Mock(side_effect=RuntimeError("inference failed"))
    source = wav()
    target = tmp_path / "existing.wav"
    target.write_bytes(b"existing output")
    with pytest.raises(RuntimeError, match="inference failed"):
        run(source, target, 0.85)
    assert target.read_bytes() == b"existing output"
    assert not list(tmp_path.glob(".voiceforge-ai-*"))