Files
ai-agent-book/chapter8/orpheus/test_orpheus_inference.py
T
liqiang b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
2026-08-20 13:12:50 +00:00

86 lines
2.6 KiB
Python

import importlib.util
import sys
from pathlib import Path
from types import ModuleType
from unittest.mock import MagicMock, patch
import pytest
import torch
def _optional_dependency_stubs():
torchaudio = ModuleType("torchaudio")
torchaudio.__path__ = []
transforms = ModuleType("torchaudio.transforms")
torchaudio.transforms = transforms
unsloth = ModuleType("unsloth")
unsloth.FastLanguageModel = MagicMock()
snac = ModuleType("snac")
snac.SNAC = MagicMock()
return {
"torchaudio": torchaudio,
"torchaudio.transforms": transforms,
"unsloth": unsloth,
"snac": snac,
}
OPTIONAL_DEPENDENCY_STUBS = _optional_dependency_stubs()
INFERENCE_PATH = Path(__file__).with_name("inference.py")
SPEC = importlib.util.spec_from_file_location("orpheus_inference_under_test", INFERENCE_PATH)
INFERENCE_MODULE = importlib.util.module_from_spec(SPEC)
# Keep heavyweight optional dependencies local to this import. patch.dict
# restores every prior sys.modules entry immediately after inference.py loads.
with patch.dict(sys.modules, OPTIONAL_DEPENDENCY_STUBS):
SPEC.loader.exec_module(INFERENCE_MODULE)
OrpheusInference = INFERENCE_MODULE.OrpheusInference
class DummyInference(OrpheusInference):
def __init__(self):
self.snac_model = MagicMock()
@pytest.mark.parametrize("tail_length", range(7))
def test_redistribute_codes_discards_trailing_incomplete_frame(tail_length):
dummy = DummyInference()
expected_audio = torch.ones(1, 1, 4)
dummy.snac_model.decode.return_value = expected_audio
# One valid SNAC frame followed by zero to six incomplete-frame codes.
valid_frame = [1, 4098, 8195, 12292, 16389, 20486, 24583]
audio = dummy._redistribute_codes(valid_frame + [999] * tail_length)
assert audio is expected_audio
dummy.snac_model.decode.assert_called_once()
codes = dummy.snac_model.decode.call_args.args[0]
assert [tensor.tolist() for tensor in codes] == [
[[1]],
[[2, 5]],
[[3, 4, 6, 7]],
]
@pytest.mark.parametrize("incomplete_length", range(1, 7))
def test_redistribute_codes_returns_silence_for_only_incomplete_codes(incomplete_length):
dummy = DummyInference()
audio = dummy._redistribute_codes([999] * incomplete_length)
assert tuple(audio.shape) == (1, 1, 1000)
assert torch.count_nonzero(audio).item() == 0
dummy.snac_model.decode.assert_not_called()
@pytest.mark.parametrize(
("module_name", "stub"),
OPTIONAL_DEPENDENCY_STUBS.items(),
ids=OPTIONAL_DEPENDENCY_STUBS,
)
def test_optional_dependency_stubs_are_restored(module_name, stub):
assert sys.modules.get(module_name) is not stub