feat: chunking, zipenhancer denoising

1 files changed, 221 insertions, 0 deletions

diff --git a/lib/tests/test_updates.py b/lib/tests/test_updates.py
new file mode 100644
index 0000000..7ad1727
--- /dev/null
+++ b/lib/tests/test_updates.py
@@ -0,0 +1,221 @@
+import subprocess
+import sys
+from types import SimpleNamespace
+
+from producer import lazy, updates
+
+
+def test_parse_would_install():
+    out = "\n".join(
+        [
+            "Using Python 3.11.16 environment at: /x",
+            "Resolved 2 packages in 1ms",
+            "Would install 2 packages",
+            " + scipy==1.14.1",
+            " - numpy==1.26.4",
+            " + numpy==2.2.6",
+            " ~ cffi==2.1.1",
+            "",
+        ]
+    )
+    assert updates._parse_would_install(out) == {
+        "scipy": "1.14.1",
+        "numpy": "2.2.6",
+        "cffi": "2.1.1",
+    }
+
+
+def test_vkey_ordering():
+    assert updates._vkey("1.26.4") < updates._vkey("2.2.6")
+    assert updates._vkey("0.5.6") < updates._vkey("0.5.7")
+    assert updates._vkey("2.7.1") < updates._vkey("2.7.1+cu126")
+    assert updates._vkey("1.2.10") > updates._vkey("1.2.9")
+    assert not updates._vkey("0.5.6") > updates._vkey("0.5.6")
+
+
+def test_core_names(tmp_path, monkeypatch):
+    reqs = tmp_path / "requirements-core.txt"
+    reqs.write_text("numpy==1.26.4\nscipy>=1.11,<1.15\n\n# comment\nsoundfile>=0.12,<0.13\n")
+    monkeypatch.setattr(updates, "CORE_FILE", reqs)
+    assert updates._core_names() == ["numpy", "scipy", "soundfile"]
+
+
+def test_probe_filters_to_keep(monkeypatch):
+    calls = []
+
+    def fake_run(cmd, capture_output, text, timeout):
+        calls.append(cmd)
+        return SimpleNamespace(
+            returncode=0,
+            stdout="Resolved 2 packages\n + scipy==1.14.1\n + numpy==2.2.6\n",
+        )
+
+    monkeypatch.setattr(lazy, "find_uv", lambda: "uv")
+    monkeypatch.setattr(updates.subprocess, "run", fake_run)
+    got = updates._probe(["-U", "-r", "reqs.txt"], None, {"scipy"})
+    assert got == {"scipy": "1.14.1"}
+    assert calls[0][1:4] == ["pip", "install", "--dry-run"]
+
+
+def test_probe_failure_returns_empty(monkeypatch):
+    def boom(cmd, capture_output, text, timeout):
+        raise subprocess.TimeoutExpired(cmd, timeout)
+
+    monkeypatch.setattr(lazy, "find_uv", lambda: "uv")
+    monkeypatch.setattr(updates.subprocess, "run", boom)
+    assert updates._probe(["torch"], None, {"torch"}) == {}
+
+
+def test_collect_full_matrix(monkeypatch, tmp_path):
+    installed = {
+        "numpy": "1.26.4",
+        "scipy": "1.13.1",
+        "soundfile": "0.12.1",
+        "pyloudnorm": "0.2.0",
+        "torch": "2.7.1+cu126",
+        "torchaudio": "2.7.1+cu126",
+        "deepfilternet": "0.5.6",
+        "zipenhancer": None,
+        "clearvoice": None,
+    }
+
+    def fake_probe(args, python, keep):
+        if "torch" in args and "torchaudio" in args and "--index-url" in args:
+            return {"torch": "2.9.0+cu126", "torchaudio": "2.9.0+cu126"}
+        if "deepfilternet" in args:
+            return {"deepfilternet": "0.5.7", "torch": "2.14.0"}
+        if "zipenhancer" in args:
+            return {}
+        return {"scipy": "1.14.1"}
+
+    monkeypatch.setattr(updates, "_installed", lambda d: installed.get(d))
+    monkeypatch.setattr(updates, "_probe", fake_probe)
+    monkeypatch.setattr(lazy, "gpu_present", lambda: True)
+    monkeypatch.setattr(lazy, "find_uv", lambda: "uv")
+    monkeypatch.setattr(updates, "RESEMBLE_PY", tmp_path / "missing" / "python")
+
+    got = updates.collect()
+    labels = [u.label for u in got]
+    assert labels == ["scipy", "torch + torchaudio", "deepfilternet"]
+
+    scipy = got[0]
+    assert (scipy.old, scipy.new) == ("1.13.1", "1.14.1")
+    assert "-U" in scipy.cmd and str(updates.CORE_FILE) in scipy.cmd
+
+    torch_u = got[1]
+    assert (torch_u.old, torch_u.new) == ("2.7.1+cu126", "2.9.0+cu126")
+    assert "CUDA" in torch_u.note
+    assert "torch==2.9.0+cu126" in torch_u.cmd
+    assert "torchaudio==2.9.0+cu126" in torch_u.cmd
+    assert lazy.TORCH_GPU_INDEX in torch_u.cmd
+
+    dfn = got[2]
+    assert (dfn.old, dfn.new) == ("0.5.6", "0.5.7")
+    assert "torch==2.7.1" in dfn.cmd
+    assert "deepfilternet==0.5.7" in dfn.cmd
+    assert "torch==2.14.0" not in dfn.cmd
+
+
+def test_collect_skips_uninstalled_engines_and_missing_torch(monkeypatch, tmp_path):
+    monkeypatch.setattr(updates, "_installed", lambda d: None)
+    monkeypatch.setattr(lazy, "find_uv", lambda: "uv")
+    monkeypatch.setattr(updates, "RESEMBLE_PY", tmp_path / "missing" / "python")
+
+    probed = []
+    monkeypatch.setattr(updates, "_probe", lambda args, python, keep: probed.append(args) or {})
+    assert updates.collect() == []
+    assert len(probed) == 1
+    assert "deepfilternet" not in probed[0] and "torch" not in probed[0]
+
+
+def test_collect_resemble_isolated_venv(monkeypatch, tmp_path):
+    py = tmp_path / "resemble" / "bin" / "python"
+    py.parent.mkdir(parents=True)
+    py.write_text("")
+    monkeypatch.setattr(updates, "_installed", lambda d: "2.7.1+cu126" if d == "torch" else None)
+    monkeypatch.setattr(updates, "_installed_in", lambda p, d: "0.0.1")
+    monkeypatch.setattr(lazy, "find_uv", lambda: "uv")
+    monkeypatch.setattr(updates, "RESEMBLE_PY", py)
+
+    def fake_probe(args, python, keep):
+        if "resemble-enhance" in args:
+            assert python == str(py)
+            return {"resemble-enhance": "0.0.2"}
+        return {}
+
+    monkeypatch.setattr(updates, "_probe", fake_probe)
+    got = updates.collect()
+    assert len(got) == 1
+    u = got[0]
+    assert u.label == "resemble-enhance"
+    assert "resemble-enhance==0.0.2" in u.cmd
+    assert str(py) in u.cmd
+
+
+def test_confirm(monkeypatch):
+    answers = iter(["", "no", "y", "yes"])
+    monkeypatch.setattr("builtins.input", lambda prompt: next(answers))
+    assert not updates._confirm("install updates?")
+    assert not updates._confirm("install updates?")
+    assert updates._confirm("install updates?")
+    assert updates._confirm("install updates?")
+
+
+def test_check_and_prompt_non_tty_notices_only(monkeypatch, capsys):
+    monkeypatch.setattr(
+        updates,
+        "collect",
+        lambda: [updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv"])],
+    )
+    monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: False))
+    assert updates.check_and_prompt() is False
+    out = capsys.readouterr().out
+    assert "updates available" in out
+    assert "./producer update" in out
+
+
+def test_check_and_prompt_assume_yes_applies(monkeypatch, capsys):
+    u = updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv", "pip", "install", "scipy"])
+    monkeypatch.setattr(updates, "collect", lambda: [u])
+    runs = []
+    monkeypatch.setattr(updates.ui, "run", lambda cmd, label, check=True: runs.append(cmd))
+    assert updates.check_and_prompt(assume_yes=True) is True
+    assert runs == [["uv", "pip", "install", "scipy"]]
+    assert "updates installed" in capsys.readouterr().out
+
+
+def test_check_and_prompt_declined(monkeypatch, capsys):
+    u = updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv"])
+    monkeypatch.setattr(updates, "collect", lambda: [u])
+    monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: True))
+    monkeypatch.setattr(updates, "_confirm", lambda prompt: False)
+    runs = []
+    monkeypatch.setattr(updates.ui, "run", lambda cmd, label, check=True: runs.append(cmd))
+    assert updates.check_and_prompt() is False
+    assert runs == []
+    assert "skipped updates" in capsys.readouterr().out
+
+
+def test_check_and_prompt_up_to_date(monkeypatch, capsys):
+    monkeypatch.setattr(updates, "collect", lambda: [])
+    assert updates.check_and_prompt(force=True) is False
+    assert "up to date" in capsys.readouterr().out
+
+
+def test_update_command_requires_yes_when_not_interactive(monkeypatch):
+    monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: False))
+    monkeypatch.setattr(
+        updates,
+        "collect",
+        lambda: [updates.Update("scipy", "1.13.1", "1.14.1", "", ["uv"])],
+    )
+    try:
+        updates.run_update_command([])
+    except SystemExit as e:
+        assert "--yes" in str(e)
+    else:
+        raise AssertionError("expected SystemExit")
+    runs = []
+    monkeypatch.setattr(updates.ui, "run", lambda cmd, label, check=True: runs.append(cmd))
+    assert updates.run_update_command(["--yes"]) == 0
+    assert len(runs) == 1