endman100-ComfyUI-WhisperLargeV3-Repack/tests/test_nodes.py

import json
import importlib.util
import sys
import types
from pathlib import Path

import pytest


PACKAGE_ROOT = Path(__file__).resolve().parents[1]
if str(PACKAGE_ROOT) not in sys.path:
    sys.path.insert(0, str(PACKAGE_ROOT))


class FakeSegment:
    id = 0
    start = 0.0
    end = 1.25
    text = " Hello world."
    words = None


class FakeInfo:
    language = "en"
    duration = 1.25


class FakeWhisperModel:
    init_args = None
    init_calls = []
    transcribe_args = None
    fail_first_transcribe = False

    def __init__(self, *args, **kwargs):
        FakeWhisperModel.init_args = (args, kwargs)
        FakeWhisperModel.init_calls.append((args, kwargs))

    def transcribe(self, *args, **kwargs):
        if FakeWhisperModel.fail_first_transcribe:
            FakeWhisperModel.fail_first_transcribe = False
            raise RuntimeError("Library cublas64_12.dll is not found or cannot be loaded")
        FakeWhisperModel.transcribe_args = (args, kwargs)
        return iter([FakeSegment()]), FakeInfo()


@pytest.fixture(autouse=True)
def fake_folder_paths(monkeypatch, tmp_path):
    module = types.ModuleType("folder_paths")
    module.models_dir = str(tmp_path / "models")
    module.get_input_directory = lambda: str(tmp_path)
    module.filter_files_content_types = lambda files, content_types: files
    module.get_annotated_filepath = lambda name: str(tmp_path / name)
    module.exists_annotated_filepath = lambda name: (tmp_path / name).exists()
    monkeypatch.setitem(sys.modules, "folder_paths", module)


@pytest.fixture(autouse=True)
def fake_faster_whisper(monkeypatch):
    module = types.ModuleType("faster_whisper")
    module.WhisperModel = FakeWhisperModel
    monkeypatch.setitem(sys.modules, "faster_whisper", module)
    yield
    FakeWhisperModel.init_args = None
    FakeWhisperModel.init_calls = []
    FakeWhisperModel.transcribe_args = None
    FakeWhisperModel.fail_first_transcribe = False


def test_transcribe_path_uses_large_v3_and_returns_segments(tmp_path):
    import nodes

    nodes._load_model.cache_clear()
    audio = tmp_path / "sample.wav"
    audio.write_bytes(b"fake wav")

    transcript, language, duration, segments_json = nodes._transcribe_path(
        str(audio),
        "large-v3",
        "cpu",
        "int8",
        "",
        "transcribe",
        5,
        True,
        False,
        "",
        0.0,
    )

    assert transcript == "Hello world."
    assert language == "en"
    assert duration == 1.25
    assert json.loads(segments_json)[0]["text"] == " Hello world."
    assert FakeWhisperModel.init_args[0][0] == "large-v3"
    assert FakeWhisperModel.transcribe_args[1]["language"] is None


def test_auto_cuda_runtime_error_falls_back_to_cpu_int8(tmp_path, monkeypatch):
    import nodes

    nodes._load_model.cache_clear()
    monkeypatch.setattr(nodes, "_resolve_device", lambda device: "cuda")
    FakeWhisperModel.fail_first_transcribe = True
    audio = tmp_path / "sample.wav"
    audio.write_bytes(b"fake wav")

    transcript, language, duration, segments_json = nodes._transcribe_path(
        str(audio),
        "large-v3",
        "auto",
        "auto",
        "",
        "transcribe",
        5,
        True,
        False,
        "",
        0.0,
    )

    assert transcript == "Hello world."
    assert language == "en"
    assert duration == 1.25
    assert json.loads(segments_json)[0]["start"] == 0.0
    assert FakeWhisperModel.init_calls[0][1]["device"] == "cuda"
    assert FakeWhisperModel.init_calls[1][1]["device"] == "cpu"
    assert FakeWhisperModel.init_calls[1][1]["compute_type"] == "int8"


def test_file_node_returns_comfyui_result(monkeypatch, tmp_path):
    import nodes

    audio = tmp_path / "sample.wav"
    audio.write_bytes(b"fake wav")
    monkeypatch.setattr(nodes.folder_paths, "get_annotated_filepath", lambda _: str(audio))
    monkeypatch.setattr(
        nodes,
        "_transcribe_path",
        lambda *args, **kwargs: ("hello", "en", 1.0, "[]"),
    )

    result = nodes.WhisperLargeV3RepackTranscribeFile().transcribe(
        "sample.wav",
        "large-v3",
        "auto",
        "auto",
        "",
        "transcribe",
        5,
        True,
        False,
        "",
        0.0,
    )

    assert result["ui"]["text"] == ("hello",)
    assert result["result"] == ("hello", "en", 1.0, "[]")


def test_package_init_loads_local_mappings_not_global_nodes(monkeypatch):
    global_nodes = types.ModuleType("nodes")
    global_nodes.NODE_CLASS_MAPPINGS = {"KSampler": object}
    global_nodes.NODE_DISPLAY_NAME_MAPPINGS = {"KSampler": "KSampler"}
    monkeypatch.setitem(sys.modules, "nodes", global_nodes)

    spec = importlib.util.spec_from_file_location(
        "whisper_repack_package_init_test",
        PACKAGE_ROOT / "__init__.py",
        submodule_search_locations=[str(PACKAGE_ROOT)],
    )
    package = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(package)

    assert "WhisperLargeV3RepackTranscribeFile" in package.NODE_CLASS_MAPPINGS
    assert "WhisperLargeV3RepackTranscribeAudio" in package.NODE_CLASS_MAPPINGS
    assert "KSampler" not in package.NODE_CLASS_MAPPINGS