"""使用模拟 VLLM 适配器验证本地 WebSocket 流程。""" from __future__ import annotations from types import SimpleNamespace import asyncio import io import wave import os from unittest.mock import AsyncMock, patch from aiohttp import web from aiohttp.test_utils import AioHTTPTestCase from server import ( AUXILIARY_SERVICE_KEY, MODEL_SERVICE_KEY, config_handler, deployment_model_name, validate_model_service_url, websocket_handler, ) class FakeModelService: native_partial_supported = False config = SimpleNamespace(base_url="http://fake/v1", model="fake-model") async def transcribe(self, audio_bytes: bytes, source: str, file_name: str, partial: bool) -> str: return "partial text" if partial else "final text" class FakeAuxiliaryService: """返回固定时间段的聚类服务，用于验证 final 后的同句 speaker 更新。""" def __init__(self) -> None: self.calls = 0 async def health(self) -> dict[str, object]: """模拟辅助模型服务已完成预加载。""" return {"ready": True, "speaker_embedding_ready": True} async def resolve_speaker( self, audio_bytes: bytes, session_id: str, start_time_ms: float, end_time_ms: float, ) -> dict[str, object]: _ = (audio_bytes, session_id, start_time_ms, end_time_ms) speaker_id = [0, 1, 0][min(self.calls, 2)] self.calls += 1 return { "speaker_id": speaker_id, "speaker_name": f"说话人 {speaker_id + 1}", "speaker_evidence": "fresh", "speaker_confidence": 0.9, "speaker_strategy": "online_embedding_cluster", } async def reset_speaker_session(self, session_id: str) -> None: _ = session_id class UnhealthyAuxiliaryService(FakeAuxiliaryService): """模拟端口可访问但辅助模型尚未就绪的服务。""" async def health(self) -> dict[str, object]: return {"ready": False, "speaker_embedding_ready": False} class WebSocketFlowTests(AioHTTPTestCase): async def collect(self, ws, until="end"): """限定等待时间，回归测试中的队列卡死必须表现为失败。""" events = [] async with asyncio.timeout(10): while True: event = await ws.receive_json() events.append(event) if event["type"] == until: return events def get_app(self) -> web.Application: app = web.Application() app[MODEL_SERVICE_KEY] = FakeModelService() app[AUXILIARY_SERVICE_KEY] = FakeAuxiliaryService() app.router.add_get("/api/config", config_handler) app.router.add_get("/ws", websocket_handler) return app def test_vllm_url_validation(self) -> None: self.assertEqual(validate_model_service_url(" http://asr.local/v1/ "), "http://asr.local/v1") with self.assertRaises(ValueError): validate_model_service_url("asr.local:8000/v1") with self.assertRaises(ValueError): validate_model_service_url("http://user:password@asr.local/v1") def test_deployment_alias_and_served_name_match_vllm(self): """WebSocket 不能把下载别名直接当成 vLLM 公开模型名。""" with patch.dict(os.environ, {"QWEN3_ASR_MODEL": "0.6b"}, clear=True): self.assertEqual(deployment_model_name(), "Qwen/Qwen3-ASR-0.6B") with patch.dict(os.environ, {"QWEN3_ASR_MODEL": "0.6b", "VLLM_SERVED_MODEL_NAME": "custom-asr"}, clear=True): self.assertEqual(deployment_model_name(), "custom-asr") async def test_empty_final_retracts_partial_instead_of_leaving_pending(self): """最终没有识别文本时撤回临时内容，不留下永远等待声纹的行。""" async def transcribe(*args, partial): return "temporary" if partial else "" self.app[MODEL_SERVICE_KEY].transcribe = transcribe ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "partial_interval_ms": 300}) await ws.receive_json() await ws.send_bytes(b"\xe8\x03" * 16000) await ws.send_json({"type": "eof"}) events = await self.collect(ws) self.assertTrue(any(e["type"] == "sentences" for e in events)) self.assertEqual(events[-1]["sentences"], []) self.assertEqual(events[-1]["display_blocks"], []) async def test_compressed_file_is_rejected_for_streaming_validation(self) -> None: ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "source": "file", "file_name": "meeting.mp3"}) error = await ws.receive_json() self.assertEqual(error["type"], "error") self.assertIn("PCM 或 WAV", error["message"]) await ws.close() async def test_short_interruption_does_not_inherit_or_call_embedding(self): """句尾静音不能凑够声纹时长，长段 A 后的短插话独立保持 pending。""" ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "source": "mic"}) await ws.receive_json() await ws.send_bytes(b"\xe8\x03" * 16000 + b"\x00\x00" * 12800 + b"\xe8\x03" * 8000 + b"\x00\x00" * 12800) await ws.send_json({"type": "stop"}) events = await self.collect(ws) final = events[-1] self.assertEqual(self.app[AUXILIARY_SERVICE_KEY].calls, 1) self.assertEqual([s["speaker_id"] for s in final["sentences"]], [0, -1]) self.assertEqual(final["sentences"][1]["speaker_status"], "insufficient_audio") self.assertEqual(final["sentences"][0]["end_time"], 1000) self.assertEqual(len(final["display_blocks"]), 2) async def test_short_speaker_gap_splits_turn_before_next_speaker(self): """说话人模式在短交接停顿处切段，保持每段独立进入在线聚类。""" ws = await self.client.ws_connect("/ws") await ws.send_json( { "type": "start", "source": "mic", "speaker_diarization": 1, "speaker_gap_ms": 400, "partial_interval_ms": 1200, } ) self.assertEqual((await ws.receive_json())["type"], "start") voiced = b"\xe8\x03" * 16000 # 每位说话人一秒有效语音 handoff_gap = b"\x00\x00" * 6400 # 400ms，低于原始 800ms 切段阈值 trailing_silence = b"\x00\x00" * 12800 await ws.send_bytes(voiced + handoff_gap + voiced + trailing_silence) await ws.send_json({"type": "eof"}) events: list[dict[str, object]] = [] while True: message = await ws.receive_json() events.append(message) if message["type"] == "end": break final_sentences = [ message["sentences"][0] for message in events if message["type"] == "sentences" and message["sentences"][0]["sentence_type"] == 1 ] latest_by_id = {int(item["sentence_id"]): item for item in final_sentences} latest = [latest_by_id[index] for index in sorted(latest_by_id)] self.assertEqual(len(latest), 2) self.assertEqual([item["speaker_id"] for item in latest], [0, 1]) self.assertEqual(self.app[AUXILIARY_SERVICE_KEY].calls, 2) await ws.close() async def test_stop_waits_for_slow_speaker_and_includes_final_snapshot(self): """在 end 前必须收到所有声纹结果，不能复现页面原先五秒断开的行为。""" auxiliary = self.app[AUXILIARY_SERVICE_KEY] original = auxiliary.resolve_speaker async def delayed(*args): await asyncio.sleep(5.1) return await original(*args) auxiliary.resolve_speaker = delayed auxiliary.reset_speaker_session = AsyncMock() ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start"}) start = await ws.receive_json() await ws.send_bytes(b"\xe8\x03" * 16000) await ws.send_json({"type": "stop"}) events = await self.collect(ws) self.assertEqual(events[-1]["sentences"][0]["speaker_id"], 0) self.assertTrue(any(e["type"] == "draining" for e in events)) self.assertTrue(any(e["type"] == "sentences" and e["sentences"][0]["speaker_status"] == "processing" for e in events)) await ws.receive() # 等待服务端执行 finally 并关闭连接 auxiliary.reset_speaker_session.assert_awaited_once_with(start["session_id"]) async def test_speaker_error_is_visible_on_segment_and_asr_finishes(self): """声纹推理失败不能吞掉转写，且每条失败片段要携带诊断原因。""" self.app[AUXILIARY_SERVICE_KEY].resolve_speaker = AsyncMock(side_effect=RuntimeError("embedding model missing")) ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start"}) await ws.receive_json() await ws.send_bytes(b"\xe8\x03" * 16000) await ws.send_json({"type": "eof"}) events = await self.collect(ws) segment = events[-1]["sentences"][0] self.assertEqual(segment["sentence"], "final text") self.assertEqual(segment["speaker_status"], "service_error") self.assertIn("embedding model missing", segment["speaker_reason"]) self.assertTrue(any(e["type"] == "speaker_warning" for e in events)) async def test_asr_failure_is_reported_before_stop(self): """音频 worker 抛异常时，接收任务应立即报告，不能等到客户端 stop。""" self.app[MODEL_SERVICE_KEY].transcribe = AsyncMock(side_effect=RuntimeError("vllm unavailable")) ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "partial_interval_ms": 300}) await ws.receive_json() await ws.send_bytes(b"\xe8\x03" * 16000) events = await self.collect(ws, until="error") self.assertIn("vllm unavailable", events[-1]["message"]) async def test_unknown_speaker_response_is_diagnosable(self): """旧服务只回标签、缺少 fresh/confidence 时应说明拒绝原因。""" self.app[AUXILIARY_SERVICE_KEY].resolve_speaker = AsyncMock(return_value={"speaker_id": 0}) ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start"}) await ws.receive_json() await ws.send_bytes(b"\xe8\x03" * 16000) await ws.send_json({"type": "eof"}) events = await self.collect(ws) self.assertEqual(events[-1]["sentences"][0]["speaker_status"], "evidence_rejected") async def test_abort_and_disconnect_release_cluster_state(self): """清理不应只存在于成功 stop 的路径。""" auxiliary = self.app[AUXILIARY_SERVICE_KEY] auxiliary.reset_speaker_session = AsyncMock() for abort in (True, False): ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start"}) await ws.receive_json() if abort: await ws.send_json({"type": "abort"}) await ws.receive() else: await ws.close() async with asyncio.timeout(2): while auxiliary.reset_speaker_session.await_count < 2: await asyncio.sleep(0.01) self.assertEqual(auxiliary.reset_speaker_session.await_count, 2) async def test_extended_wav_header_is_removed_before_asr(self): """分片 RIFF/JUNK/fmt/data 头不能混入声纹和 ASR 的 PCM 数据。""" output = io.BytesIO() pcm = b"\xe8\x03" * 16000 with wave.open(output, "wb") as wav: wav.setparams((1, 2, 16000, 0, "NONE", "")) wav.writeframes(pcm) original = output.getvalue() junk = b"JUNK\x04\x00\x00\x00test" payload = b"RIFF" + (len(original) - 8 + len(junk)).to_bytes(4, "little") + original[8:12] + junk + original[12:] transcribe = AsyncMock(return_value="wav text") self.app[MODEL_SERVICE_KEY].transcribe = transcribe ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "source": "file", "file_name": "test.wav"}) await ws.receive_json() for offset in range(0, len(payload), 337): await ws.send_bytes(payload[offset:offset + 337]) await ws.send_json({"type": "eof"}) events = await self.collect(ws) self.assertEqual(transcribe.call_args.args[0], pcm) self.assertEqual(events[-1]["sentences"][0]["end_time"], 1000) async def test_incompatible_wav_is_rejected_immediately(self): """非 16kHz 单声道 WAV 不能被误解释为可识别的 PCM16。""" output = io.BytesIO() with wave.open(output, "wb") as wav: wav.setparams((2, 2, 44100, 0, "NONE", "")) wav.writeframes(b"\x00\x00" * 2000) ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "source": "file", "file_name": "bad.wav"}) await ws.receive_json() await ws.send_bytes(output.getvalue()) events = await self.collect(ws, until="error") self.assertIn("16kHz", events[-1]["message"]) async def test_frontend_can_read_default_vllm_config(self) -> None: response = await self.client.get("/api/config") self.assertEqual(response.status, 200) self.assertEqual(await response.json(), {"model_service_url": "http://fake/v1", "model": "fake-model", "speaker_service_url": None}) async def test_partial_and_final_share_one_sentence_id(self) -> None: ws = await self.client.ws_connect("/ws") await ws.send_json( { "type": "start", "source": "mic", "speaker_diarization": 0, "partial_interval_ms": 300, "max_segment_sec": 12, } ) start = await ws.receive_json() self.assertEqual(start["type"], "start") # 使用幅度足够的 PCM 语音帧，静音帧会被新切句器正确忽略。 await ws.send_bytes(b"\xe8\x03" * 16000) messages = [] while True: message = await ws.receive_json() messages.append(message) if message["type"] == "sentences": self.assertEqual(message["sentences"][0]["sentence_id"], 0) if message["sentences"][0]["sentence_type"] == 0: break await ws.send_json({"type": "eof"}) while True: message = await ws.receive_json() messages.append(message) if message["type"] == "end": break sentence_events = [message for message in messages if message["type"] == "sentences"] self.assertGreaterEqual(len(sentence_events), 2) self.assertTrue(all(event["sentences"][0]["sentence_id"] == 0 for event in sentence_events)) self.assertTrue(all(event["sentences"][0]["sentence_type"] == 0 for event in sentence_events[:-1])) self.assertEqual(sentence_events[-1]["sentences"][0]["sentence_type"], 1) self.assertEqual(sentence_events[-1]["sentences"][0]["sentence"], "final text") await ws.close() async def test_speaker_health_failure_is_reported_without_blocking_asr(self) -> None: """辅助模型未就绪时先报告告警，同时保留 ASR 会话能力。""" self.app[AUXILIARY_SERVICE_KEY] = UnhealthyAuxiliaryService() ws = await self.client.ws_connect("/ws") await ws.send_json({"type": "start", "source": "mic", "speaker_diarization": 1}) start = await ws.receive_json() warning = await ws.receive_json() self.assertEqual(start["type"], "start") self.assertFalse(start["speaker_service_health"]["ready"]) self.assertFalse(start["speaker_gap_enabled"]) self.assertEqual(warning["type"], "speaker_warning") self.assertIn("未就绪", warning["message"]) await ws.close() async def test_vad_split_and_speaker_update(self) -> None: ws = await self.client.ws_connect("/ws") await ws.send_json( { "type": "start", "source": "mic", "speaker_diarization": 1, "partial_interval_ms": 1200, } ) self.assertEqual((await ws.receive_json())["type"], "start") voiced = b"\xe8\x03" * 16000 # 每段一秒有效语音，满足独立声纹长度要求 silence = b"\x00\x00" * 12800 # 0.8 秒静音，触发当前 turn 提交 await ws.send_bytes(voiced + silence + voiced + silence + voiced) await ws.send_json({"type": "eof"}) events: list[dict[str, object]] = [] while True: message = await ws.receive_json() events.append(message) if message["type"] == "end": break sentence_events = [message for message in events if message["type"] == "sentences"] final_sentences = [ message["sentences"][0] for message in sentence_events if message["sentences"][0]["sentence_type"] == 1 ] latest_by_id = {int(item["sentence_id"]): item for item in final_sentences} self.assertGreaterEqual(len(latest_by_id), 2) latest = [latest_by_id[index] for index in sorted(latest_by_id)[-3:]] self.assertEqual([item["speaker_id"] for item in latest], [0, 1, 0]) self.assertTrue(all(item["speaker_evidence"] == "fresh" for item in latest)) await ws.close()