feat: chunking, zipenhancer denoising

1 files changed, 169 insertions, 0 deletions

diff --git a/lib/tests/test_ui.py b/lib/tests/test_ui.py
new file mode 100644
index 0000000..156b3e8
--- /dev/null
+++ b/lib/tests/test_ui.py
@@ -0,0 +1,169 @@
+import hashlib
+import subprocess
+import sys
+
+import pytest
+
+from producer import ui
+
+
+class _FakeResp:
+    def __init__(self, payload: bytes, length: str | None = None):
+        self._payload = payload
+        self.headers = {"Content-Length": length} if length else {}
+
+    def read(self, n=-1):
+        if not self._payload:
+            return b""
+        out, self._payload = self._payload[:n], self._payload[n:]
+        return out
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_a):
+        return False
+
+
+def test_fmt_bytes():
+    assert ui.fmt_bytes(512) == "512 B"
+    assert ui.fmt_bytes(2048) == "2.0 KB"
+    assert ui.fmt_bytes(8 << 20) == "8.0 MB"
+    assert ui.fmt_bytes(None) == "?"
+
+
+def test_fmt_secs():
+    assert ui.fmt_secs(0) == "0:00"
+    assert ui.fmt_secs(59) == "0:59"
+    assert ui.fmt_secs(61) == "1:01"
+    assert ui.fmt_secs(3700) == "1:01:40"
+    assert ui.fmt_secs(None) == "?"
+
+
+def test_bar_fills():
+    assert ui._bar(0.0) == "[" + "-" * 22 + "]"
+    assert ui._bar(1.0) == "[" + "#" * 22 + "]"
+    half = ui._bar(0.5)
+    assert half.count("#") == 11 and half.count("-") == 11
+
+
+def test_download_writes_file(tmp_path, monkeypatch, capsys):
+    payload = b"x" * (2 << 20)
+    sha = hashlib.sha256(payload).hexdigest()
+    monkeypatch.setattr(
+        "urllib.request.urlopen", lambda req, timeout=None: _FakeResp(payload, str(len(payload)))
+    )
+    dest = tmp_path / "big.bin"
+    out = ui.download("https://example.com/big.bin", dest, sha256=sha)
+    assert out == dest
+    assert dest.read_bytes() == payload
+    assert "done" in capsys.readouterr().out
+
+
+def test_download_checksum_mismatch(tmp_path, monkeypatch):
+    monkeypatch.setattr("urllib.request.urlopen", lambda req, timeout=None: _FakeResp(b"abc"))
+    dest = tmp_path / "f.bin"
+    with pytest.raises(RuntimeError):
+        ui.download("https://example.com/f.bin", dest, sha256="0" * 64)
+    assert not dest.exists()
+    assert not dest.with_name(dest.name + ".part").exists()
+
+
+def test_download_cached_skips(tmp_path, monkeypatch, capsys):
+    dest = tmp_path / "cached.bin"
+    dest.write_bytes(b"hello")
+
+    def boom(*_a, **_k):
+        raise AssertionError("should not download")
+
+    monkeypatch.setattr("urllib.request.urlopen", boom)
+    ui.download("https://example.com/cached.bin", dest)
+    assert dest.read_bytes() == b"hello"
+    assert "cached" in capsys.readouterr().out
+
+
+def test_download_size_mismatch_redownloads(tmp_path, monkeypatch):
+    dest = tmp_path / "short.bin"
+    dest.write_bytes(b"too short")
+    monkeypatch.setattr("urllib.request.urlopen", lambda req, timeout=None: _FakeResp(b"abcd"))
+    ui.download("https://example.com/short.bin", dest, expected_size=4)
+    assert dest.read_bytes() == b"abcd"
+
+
+def test_download_milestones_non_tty(tmp_path, monkeypatch, capsys):
+    payload = b"y" * (8 << 20)
+    monkeypatch.setattr(
+        "urllib.request.urlopen", lambda req, timeout=None: _FakeResp(payload, str(len(payload)))
+    )
+    ui.download("https://example.com/m.bin", tmp_path / "m.bin")
+    out = capsys.readouterr().out
+    assert "8.0 MB/8.0 MB" in out
+    assert "done" in out
+
+
+def test_run_piped_forwards_output(monkeypatch, capsys):
+    monkeypatch.setattr(ui, "is_tty", lambda: False)
+    rc = ui.run([sys.executable, "-c", "print('hello-uv-output')"], "test run")
+    assert rc == 0
+    assert "hello-uv-output" in capsys.readouterr().out
+
+
+def test_run_piped_raises_on_failure(monkeypatch):
+    monkeypatch.setattr(ui, "is_tty", lambda: False)
+    with pytest.raises(subprocess.CalledProcessError):
+        ui.run([sys.executable, "-c", "raise SystemExit(3)"], "test run")
+
+
+def test_status_non_tty_milestones_and_done(capsys):
+    s = ui.Status(prefix="[producer] in.wav — ")
+    s.stage("denoise")
+    for i in range(1, 11):
+        s.tick(i, 10)
+    s.stage_done("denoise", 2.5)
+    s.stage("dsp")
+    s.stage_done("dsp", 0.3)
+    s.finish()
+    out = capsys.readouterr().out
+    assert "denoise 10/10 chunks (100%)" in out
+    assert "denoise done in 2.5s (10/10 chunks)" in out
+    assert "dsp done in 0.3s" in out
+
+
+def test_status_tty_live_line(monkeypatch, capsys):
+    monkeypatch.setattr(ui, "is_tty", lambda: True)
+    s = ui.Status(prefix="p — ")
+    s.stage("denoise")
+    s.tick(10, 10)  # completion redraws immediately (bypasses the 0.1 s throttle)
+    out = capsys.readouterr().out
+    assert "denoise" in out and "10/10" in out
+    s.finish()
+
+
+def test_progress_milestones_non_tty(capsys):
+    p = ui.Progress("downloading x", total=1000)
+    for i in (100, 300, 1000):
+        p.update(i)
+    p.close()
+    out = capsys.readouterr().out
+    assert "1000 B/1000 B" in out
+    assert "done (1000 B" in out
+
+
+def test_progress_tty_bar(monkeypatch, capsys):
+    monkeypatch.setattr(ui, "is_tty", lambda: True)
+    p = ui.Progress("downloading torch", total=100)
+    p.update(50)
+    p.close()
+    out = capsys.readouterr().out
+    assert "50 B/100 B (50%)" in out
+
+
+def test_log_clears_live_line(monkeypatch, capsys):
+    monkeypatch.setattr(ui, "is_tty", lambda: True)
+    p = ui.Progress("downloading", total=100)
+    p.update(50)
+    ui.log("[producer] some message")
+    out = capsys.readouterr().out
+    assert "some message" in out
+    assert out.index("some message") > out.index("downloading")
+    p.close()