98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
# Integration test for the Multi Audio Loader nodes — part of ComfyUI-AI2Go-Utils. GPL-3.0.
|
|
import json
|
|
import wave
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("comfy_api")
|
|
pytest.importorskip("av")
|
|
|
|
|
|
@pytest.fixture
|
|
def input_dir(tmp_path, monkeypatch):
|
|
import folder_paths
|
|
monkeypatch.setattr(folder_paths, "get_input_directory", lambda: str(tmp_path))
|
|
return tmp_path
|
|
|
|
|
|
def _write_wav(path, sample_rate=22050, n=2205):
|
|
with wave.open(str(path), "wb") as w:
|
|
w.setnchannels(1)
|
|
w.setsampwidth(2) # 16-bit PCM
|
|
w.setframerate(sample_rate)
|
|
w.writeframes(b"\x00\x10" * n)
|
|
|
|
|
|
def test_advanced_audio_and_filenames(input_dir):
|
|
import torch
|
|
from nodes.multi_audio_loader import _run_audio
|
|
_write_wav(input_dir / "vo.wav", sample_rate=22050)
|
|
_write_wav(input_dir / "bed.wav", sample_rate=44100)
|
|
out = _run_audio(json.dumps([{"name": "vo.wav"}, {"name": "bed.wav"}]), advanced=True)
|
|
|
|
assert len(out) == 17 and out[0] == 2
|
|
a1, name1, a2 = out[1], out[2], out[3]
|
|
assert a1["sample_rate"] == 22050 # native rate, no resampling
|
|
assert a2["sample_rate"] == 44100
|
|
assert a1["waveform"].shape == (1, 1, 2205) # batch=1, mono, 2205 samples
|
|
assert a1["waveform"].dtype == torch.float32
|
|
assert name1 == "vo.wav"
|
|
assert out[5] is None # padding after file 2's group
|
|
|
|
|
|
def test_simple_layout(input_dir):
|
|
from nodes.multi_audio_loader import _run_audio
|
|
_write_wav(input_dir / "one.wav")
|
|
out = _run_audio(json.dumps([{"name": "one.wav"}]), advanced=False)
|
|
assert len(out) == 9 and out[0] == 1
|
|
assert out[1]["sample_rate"] == 22050 and out[2] is None
|
|
|
|
|
|
def test_missing_file_raises(input_dir):
|
|
from nodes.multi_audio_loader import _run_audio
|
|
with pytest.raises(ValueError, match="gone.wav"):
|
|
_run_audio(json.dumps([{"name": "gone.wav"}]), advanced=False)
|
|
|
|
|
|
def test_corrupt_file_raises(input_dir):
|
|
from nodes.multi_audio_loader import _run_audio
|
|
# Write a text file with .wav extension
|
|
(input_dir / "bad.wav").write_text("not audio data")
|
|
with pytest.raises(ValueError, match="bad.wav"):
|
|
_run_audio(json.dumps([{"name": "bad.wav"}]), advanced=False)
|
|
|
|
|
|
def test_disabled_row_holds_position(input_dir):
|
|
# Middle file disabled: its slots stay None (not skipped-and-shifted), the third file
|
|
# still lands in its own audio_3/filename_3 group, and count only reflects the two
|
|
# enabled files.
|
|
from nodes.multi_audio_loader import _run_audio
|
|
_write_wav(input_dir / "a.wav")
|
|
_write_wav(input_dir / "b.wav")
|
|
_write_wav(input_dir / "c.wav")
|
|
out = _run_audio(json.dumps([
|
|
{"name": "a.wav"},
|
|
{"name": "b.wav", "enabled": False},
|
|
{"name": "c.wav"},
|
|
]), advanced=True)
|
|
|
|
assert out[0] == 2 # count = enabled files, not rows
|
|
assert out[1] is not None # audio_1
|
|
assert out[3] is None and out[4] is None # audio_2/filename_2 both None
|
|
audio3, name3 = out[5], out[6] # audio_3/filename_3 — its own slots, not moved up
|
|
assert audio3 is not None and name3 == "c.wav"
|
|
|
|
|
|
def test_all_disabled_yields_count_zero(input_dir):
|
|
# All rows disabled (Toggle All off): count == 0, all output slots are None, no exception.
|
|
from nodes.multi_audio_loader import _run_audio
|
|
_write_wav(input_dir / "a.wav")
|
|
_write_wav(input_dir / "b.wav")
|
|
out = _run_audio(json.dumps([
|
|
{"name": "a.wav", "enabled": False},
|
|
{"name": "b.wav", "enabled": False},
|
|
]), advanced=True)
|
|
|
|
assert out[0] == 0 # count == 0, not > 0 (no exception)
|
|
assert all(v is None for v in out[1:]) # all output slots None
|