116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
import runpy
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
from wov_node_whisper.__main__ import format_timestamp, invoke
|
|
from wov_sdk.models import InvokeRequest
|
|
|
|
|
|
class FakeSegment:
|
|
def __init__(self, start, end, text):
|
|
self.start = start
|
|
self.end = end
|
|
self.text = text
|
|
|
|
|
|
class FakeWhisperModel:
|
|
instances: list[tuple[tuple, dict]] = []
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
FakeWhisperModel.instances.append((args, kwargs))
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
|
|
def transcribe(self, path, **kwargs):
|
|
self.transcribe_args = (path, kwargs)
|
|
return (
|
|
[
|
|
FakeSegment(0, 1, "第一段"),
|
|
FakeSegment(3600.5, 3602.25, "第二段"),
|
|
],
|
|
None,
|
|
)
|
|
|
|
|
|
def _install_fake_whisper(monkeypatch, model_class=FakeWhisperModel) -> None:
|
|
fake_module = types.SimpleNamespace(WhisperModel=model_class)
|
|
monkeypatch.setitem(sys.modules, "faster_whisper", fake_module)
|
|
|
|
|
|
def _request(tmp_path, **overrides) -> InvokeRequest:
|
|
payload = {
|
|
"run_id": "run_1",
|
|
"node_instance_id": "ni_1",
|
|
"inputs": {"audio_uri": str(tmp_path / "audio.wav")},
|
|
"params": {"language": "ja"},
|
|
"output_dir": str(tmp_path / "out"),
|
|
}
|
|
payload.update(overrides)
|
|
return InvokeRequest(**payload)
|
|
|
|
|
|
def test_format_timestamp() -> None:
|
|
assert format_timestamp(0) == "00:00:00,000"
|
|
assert format_timestamp(3600.5) == "01:00:00,500"
|
|
assert format_timestamp(61.25) == "00:01:01,250"
|
|
|
|
|
|
def test_success(tmp_path, monkeypatch) -> None:
|
|
FakeWhisperModel.instances.clear()
|
|
_install_fake_whisper(monkeypatch)
|
|
(tmp_path / "audio.wav").write_bytes(b"fake")
|
|
response = invoke(_request(tmp_path))
|
|
assert response.status == "completed"
|
|
content = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8")
|
|
assert "第一段" in content
|
|
assert "01:00:00,500 --> 01:00:02,250" in content
|
|
_, kwargs = FakeWhisperModel.instances[-1]
|
|
assert kwargs["device"] == "auto"
|
|
assert kwargs["compute_type"] == "auto"
|
|
|
|
|
|
def test_compute_type_override(tmp_path, monkeypatch) -> None:
|
|
FakeWhisperModel.instances.clear()
|
|
_install_fake_whisper(monkeypatch)
|
|
(tmp_path / "audio.wav").write_bytes(b"fake")
|
|
response = invoke(_request(tmp_path, params={"language": "ja", "compute_type": "int8"}))
|
|
assert response.status == "completed"
|
|
_, kwargs = FakeWhisperModel.instances[-1]
|
|
assert kwargs["compute_type"] == "int8"
|
|
|
|
|
|
def test_model_raises(tmp_path, monkeypatch) -> None:
|
|
class BrokenModel:
|
|
def __init__(self, *args, **kwargs):
|
|
raise RuntimeError("model load failed")
|
|
|
|
_install_fake_whisper(monkeypatch, BrokenModel)
|
|
(tmp_path / "audio.wav").write_bytes(b"fake")
|
|
response = invoke(_request(tmp_path))
|
|
assert response.status == "failed"
|
|
assert "model load failed" in response.error
|
|
|
|
|
|
def test_missing_input(tmp_path) -> None:
|
|
response = invoke(_request(tmp_path, inputs={}))
|
|
assert response.status == "failed"
|
|
|
|
|
|
def test_missing_file(tmp_path) -> None:
|
|
response = invoke(_request(tmp_path))
|
|
assert response.status == "failed"
|
|
assert "audio file not found" in response.error
|
|
|
|
|
|
def test_entrypoint(monkeypatch) -> None:
|
|
module_path = Path(__file__).resolve().parent.parent / "wov_node_whisper" / "__main__.py"
|
|
captured = {}
|
|
|
|
def fake_run_node(manifest, handler) -> None:
|
|
captured["id"] = manifest.id
|
|
|
|
monkeypatch.setattr("wov_sdk.server.run_node", fake_run_node)
|
|
runpy.run_path(str(module_path), run_name="__main__")
|
|
assert captured["id"] == "faster-whisper"
|