wehub-resource-sync
118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
"""Unit tests for base STT `RecognizeStream` fields (start_time, etc.)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from livekit.agents import APIConnectionError
|
|
from livekit.agents.stt import (
|
|
STT,
|
|
RecognizeStream,
|
|
SpeechData,
|
|
SpeechEvent,
|
|
SpeechEventType,
|
|
STTCapabilities,
|
|
)
|
|
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS
|
|
from livekit.agents.utils.audio import AudioBuffer
|
|
|
|
pytestmark = [pytest.mark.unit, pytest.mark.virtual_time, pytest.mark.no_concurrent]
|
|
|
|
|
|
class _DummyStream(RecognizeStream):
|
|
"""Minimal RecognizeStream for unit tests — does not hit the network."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
stt: STT,
|
|
fail_first_run: bool = False,
|
|
) -> None:
|
|
super().__init__(stt=stt, conn_options=DEFAULT_API_CONNECT_OPTIONS)
|
|
self._fail_first_run = fail_first_run
|
|
self._run_count = 0
|
|
|
|
async def _run(self) -> None:
|
|
self._run_count += 1
|
|
if self._fail_first_run and self._run_count == 1:
|
|
raise APIConnectionError("fake failure to trigger retry")
|
|
# emit a final and exit so _main_task can complete normally
|
|
self._event_ch.send_nowait(
|
|
SpeechEvent(
|
|
type=SpeechEventType.FINAL_TRANSCRIPT,
|
|
alternatives=[SpeechData(language="", text="hello")],
|
|
)
|
|
)
|
|
|
|
|
|
class _DummySTT(STT):
|
|
def __init__(self) -> None:
|
|
super().__init__(capabilities=STTCapabilities(streaming=True, interim_results=False))
|
|
|
|
async def _recognize_impl(self, buffer: AudioBuffer, *, language, conn_options) -> SpeechEvent:
|
|
raise NotImplementedError
|
|
|
|
def stream(self, *, language=None, conn_options=DEFAULT_API_CONNECT_OPTIONS) -> _DummyStream:
|
|
return _DummyStream(stt=self)
|
|
|
|
|
|
async def test_start_time_seeded_on_init() -> None:
|
|
"""start_time is initialized to approximately time.time() when the stream is created."""
|
|
stt = _DummySTT()
|
|
before = time.time()
|
|
stream = stt.stream()
|
|
after = time.time()
|
|
|
|
assert before <= stream.start_time <= after
|
|
await stream.aclose()
|
|
|
|
|
|
async def test_start_time_setter_accepts_valid_values() -> None:
|
|
"""Plugins can override start_time by assigning to the public property."""
|
|
stt = _DummySTT()
|
|
stream = stt.stream()
|
|
|
|
new_anchor = time.time() + 10.0
|
|
stream.start_time = new_anchor
|
|
assert stream.start_time == new_anchor
|
|
|
|
await stream.aclose()
|
|
|
|
|
|
async def test_start_time_setter_rejects_negative() -> None:
|
|
"""start_time setter validates non-negative, matching start_time_offset behavior."""
|
|
stt = _DummySTT()
|
|
stream = stt.stream()
|
|
|
|
with pytest.raises(ValueError, match="start_time must be non-negative"):
|
|
stream.start_time = -1.0
|
|
|
|
await stream.aclose()
|
|
|
|
|
|
async def test_start_time_reseeded_on_retry() -> None:
|
|
"""When _main_task retries after an APIError, start_time is re-seeded so plugin
|
|
overrides from the previous connection don't leak into the new one."""
|
|
stt = _DummySTT()
|
|
stream = _DummyStream(stt=stt, fail_first_run=True)
|
|
|
|
# Simulate a plugin overriding start_time during the first (failing) _run()
|
|
# by assigning a sentinel value before the task picks up.
|
|
sentinel = 1.0
|
|
stream.start_time = sentinel
|
|
|
|
# Let the main task run: it should retry past the first-run APIError, and
|
|
# on each attempt re-seed start_time to a fresh time.time() value before
|
|
# _run() is called.
|
|
await asyncio.wait_for(stream._task, timeout=5.0)
|
|
|
|
# After the retry, start_time must have been re-seeded (not equal to sentinel).
|
|
assert stream.start_time != sentinel
|
|
# And it should be a recent wall-clock value.
|
|
assert time.time() - stream.start_time < 5.0
|
|
|
|
await stream.aclose()
|