Aaron Boxer
701 lines
25 KiB
Python
701 lines
25 KiB
Python
import json
|
|
import time
|
|
import threading
|
|
import collections
|
|
import unittest
|
|
from unittest import mock
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from whisper_live.server import TranscriptionServer, BackendType, ClientManager
|
|
|
|
|
|
class TestClientManagerAddRemove(unittest.TestCase):
|
|
def setUp(self):
|
|
self.cm = ClientManager(max_clients=2, max_connection_time=60)
|
|
|
|
def test_add_and_get_client(self):
|
|
ws = MagicMock()
|
|
client = MagicMock()
|
|
self.cm.add_client(ws, client)
|
|
self.assertIs(self.cm.get_client(ws), client)
|
|
|
|
def test_get_nonexistent_client(self):
|
|
ws = MagicMock()
|
|
self.assertFalse(self.cm.get_client(ws))
|
|
|
|
def test_remove_client_calls_cleanup(self):
|
|
ws = MagicMock()
|
|
client = MagicMock()
|
|
self.cm.add_client(ws, client)
|
|
self.cm.remove_client(ws)
|
|
client.cleanup.assert_called_once()
|
|
self.assertNotIn(ws, self.cm.clients)
|
|
self.assertNotIn(ws, self.cm.start_times)
|
|
|
|
def test_remove_nonexistent_client_no_error(self):
|
|
ws = MagicMock()
|
|
self.cm.remove_client(ws) # should not raise
|
|
|
|
|
|
class TestClientManagerThreadSafety(unittest.TestCase):
|
|
def test_concurrent_add_remove(self):
|
|
cm = ClientManager(max_clients=100, max_connection_time=600)
|
|
errors = []
|
|
|
|
def add_clients(start_idx):
|
|
try:
|
|
for i in range(50):
|
|
ws = MagicMock(name=f"ws-{start_idx}-{i}")
|
|
client = MagicMock(name=f"client-{start_idx}-{i}")
|
|
cm.add_client(ws, client)
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
def remove_clients():
|
|
try:
|
|
for _ in range(25):
|
|
with cm.lock:
|
|
if cm.clients:
|
|
ws = next(iter(cm.clients))
|
|
else:
|
|
continue
|
|
cm.remove_client(ws)
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
threads = [
|
|
threading.Thread(target=add_clients, args=(0,)),
|
|
threading.Thread(target=add_clients, args=(1,)),
|
|
threading.Thread(target=remove_clients),
|
|
threading.Thread(target=remove_clients),
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
self.assertEqual(errors, [])
|
|
|
|
def test_concurrent_get_client(self):
|
|
cm = ClientManager(max_clients=100, max_connection_time=600)
|
|
ws = MagicMock()
|
|
client = MagicMock()
|
|
cm.add_client(ws, client)
|
|
errors = []
|
|
results = []
|
|
|
|
def get_many():
|
|
try:
|
|
for _ in range(100):
|
|
results.append(cm.get_client(ws))
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=get_many) for _ in range(4)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
self.assertEqual(errors, [])
|
|
self.assertTrue(all(r is client for r in results))
|
|
|
|
|
|
class TestClientManagerServerFull(unittest.TestCase):
|
|
def setUp(self):
|
|
self.cm = ClientManager(max_clients=1, max_connection_time=60)
|
|
|
|
def test_not_full_returns_false(self):
|
|
ws = MagicMock()
|
|
options = {"uid": "test"}
|
|
self.assertFalse(self.cm.is_server_full(ws, options))
|
|
|
|
def test_full_sends_wait_and_returns_true(self):
|
|
ws1 = MagicMock()
|
|
self.cm.add_client(ws1, MagicMock())
|
|
|
|
ws2 = MagicMock()
|
|
options = {"uid": "new-client"}
|
|
self.assertTrue(self.cm.is_server_full(ws2, options))
|
|
ws2.send.assert_called_once()
|
|
sent = json.loads(ws2.send.call_args[0][0])
|
|
self.assertEqual(sent["status"], "WAIT")
|
|
self.assertEqual(sent["uid"], "new-client")
|
|
|
|
|
|
class TestClientManagerTimeout(unittest.TestCase):
|
|
def setUp(self):
|
|
self.cm = ClientManager(max_clients=4, max_connection_time=10)
|
|
|
|
def test_not_timed_out(self):
|
|
ws = MagicMock()
|
|
client = MagicMock()
|
|
self.cm.add_client(ws, client)
|
|
self.assertFalse(self.cm.is_client_timeout(ws))
|
|
|
|
def test_timed_out(self):
|
|
ws = MagicMock()
|
|
client = MagicMock()
|
|
self.cm.add_client(ws, client)
|
|
self.cm.start_times[ws] = time.time() - 20
|
|
self.assertTrue(self.cm.is_client_timeout(ws))
|
|
client.disconnect.assert_called_once()
|
|
|
|
|
|
class TestClientManagerGetWaitTime(unittest.TestCase):
|
|
def test_no_clients_returns_zero(self):
|
|
cm = ClientManager(max_clients=4, max_connection_time=600)
|
|
self.assertEqual(cm.get_wait_time(), 0)
|
|
|
|
def test_single_client_wait_time(self):
|
|
cm = ClientManager(max_clients=4, max_connection_time=600)
|
|
ws = MagicMock()
|
|
cm.add_client(ws, MagicMock())
|
|
cm.start_times[ws] = time.time() - 300
|
|
wait = cm.get_wait_time()
|
|
self.assertAlmostEqual(wait, 5.0, places=0)
|
|
|
|
def test_multiple_clients_returns_minimum(self):
|
|
cm = ClientManager(max_clients=4, max_connection_time=600)
|
|
ws1, ws2 = MagicMock(), MagicMock()
|
|
cm.add_client(ws1, MagicMock())
|
|
cm.add_client(ws2, MagicMock())
|
|
cm.start_times[ws1] = time.time() - 100
|
|
cm.start_times[ws2] = time.time() - 500
|
|
wait = cm.get_wait_time()
|
|
# ws2 has 100s remaining = ~1.67 minutes
|
|
self.assertAlmostEqual(wait, 100 / 60, places=0)
|
|
|
|
|
|
class TestBackendType(unittest.TestCase):
|
|
def test_valid_types(self):
|
|
valid = BackendType.valid_types()
|
|
self.assertIn("faster_whisper", valid)
|
|
self.assertIn("tensorrt", valid)
|
|
self.assertIn("openvino", valid)
|
|
|
|
def test_is_valid(self):
|
|
self.assertTrue(BackendType.is_valid("faster_whisper"))
|
|
self.assertFalse(BackendType.is_valid("nonexistent"))
|
|
|
|
def test_type_checks(self):
|
|
self.assertTrue(BackendType.FASTER_WHISPER.is_faster_whisper())
|
|
self.assertFalse(BackendType.FASTER_WHISPER.is_tensorrt())
|
|
self.assertTrue(BackendType.TENSORRT.is_tensorrt())
|
|
self.assertTrue(BackendType.OPENVINO.is_openvino())
|
|
|
|
def test_enum_from_string(self):
|
|
bt = BackendType("faster_whisper")
|
|
self.assertEqual(bt, BackendType.FASTER_WHISPER)
|
|
|
|
def test_invalid_enum_raises(self):
|
|
with self.assertRaises(ValueError):
|
|
BackendType("invalid_backend")
|
|
|
|
|
|
class TestTranscriptionServerInit(unittest.TestCase):
|
|
def test_defaults(self):
|
|
server = TranscriptionServer()
|
|
self.assertIsNone(server.client_manager)
|
|
self.assertTrue(server.use_vad)
|
|
self.assertFalse(server.single_model)
|
|
self.assertIsNone(server.batch_config)
|
|
|
|
def test_run_invalid_backend_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(host="localhost", port=9090, backend="nonexistent")
|
|
|
|
def test_run_invalid_trt_path_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(
|
|
host="localhost",
|
|
port=9090,
|
|
backend="tensorrt",
|
|
whisper_tensorrt_path="/nonexistent/path",
|
|
)
|
|
|
|
def test_run_max_clients_zero_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(host="localhost", port=9090, max_clients=0)
|
|
|
|
def test_run_max_clients_negative_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(host="localhost", port=9090, max_clients=-1)
|
|
|
|
def test_run_max_connection_time_zero_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(host="localhost", port=9090, max_connection_time=0)
|
|
|
|
def test_run_batch_max_size_zero_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(host="localhost", port=9090, batch_enabled=True, batch_max_size=0)
|
|
|
|
def test_run_batch_window_ms_negative_raises(self):
|
|
server = TranscriptionServer()
|
|
with self.assertRaises(ValueError):
|
|
server.run(host="localhost", port=9090, batch_enabled=True, batch_window_ms=-1)
|
|
|
|
|
|
class TestTranscriptionServerGetAudio(unittest.TestCase):
|
|
def setUp(self):
|
|
self.server = TranscriptionServer()
|
|
|
|
def test_end_of_audio_returns_false(self):
|
|
ws = MagicMock()
|
|
ws.recv.return_value = b"END_OF_AUDIO"
|
|
result = self.server.get_audio_from_websocket(ws)
|
|
self.assertFalse(result)
|
|
|
|
def test_valid_audio_returns_numpy(self):
|
|
import numpy as np
|
|
ws = MagicMock()
|
|
audio = np.array([0.1, 0.2, 0.3], dtype=np.float32)
|
|
ws.recv.return_value = audio.tobytes()
|
|
result = self.server.get_audio_from_websocket(ws)
|
|
np.testing.assert_array_almost_equal(result, audio)
|
|
|
|
def test_raw_pcm_input_normalizes_int16(self):
|
|
import numpy as np
|
|
self.server.raw_pcm_input = True
|
|
ws = MagicMock()
|
|
pcm = np.array([0, 16384, -16384, 32767], dtype=np.int16)
|
|
ws.recv.return_value = pcm.tobytes()
|
|
result = self.server.get_audio_from_websocket(ws)
|
|
expected = pcm.astype(np.float32) / 32768.0
|
|
np.testing.assert_array_almost_equal(result, expected)
|
|
self.assertTrue(result.dtype == np.float32)
|
|
self.assertTrue(np.all(result >= -1.0))
|
|
self.assertTrue(np.all(result <= 1.0))
|
|
|
|
def test_uint8_audio_format_normalizes_unsigned_pcm(self):
|
|
import numpy as np
|
|
ws = MagicMock()
|
|
self.server.audio_formats[ws] = "uint8"
|
|
pcm = np.array([0, 128, 255], dtype=np.uint8)
|
|
ws.recv.return_value = pcm.tobytes()
|
|
result = self.server.get_audio_from_websocket(ws)
|
|
expected = (pcm.astype(np.float32) - 128.0) / 128.0
|
|
np.testing.assert_array_almost_equal(result, expected)
|
|
|
|
def test_raw_pcm_input_off_reads_float32(self):
|
|
import numpy as np
|
|
self.server.raw_pcm_input = False
|
|
ws = MagicMock()
|
|
audio = np.array([0.5, -0.5], dtype=np.float32)
|
|
ws.recv.return_value = audio.tobytes()
|
|
result = self.server.get_audio_from_websocket(ws)
|
|
np.testing.assert_array_almost_equal(result, audio)
|
|
|
|
|
|
class TestTranscriptionServerHandleNewConnection(unittest.TestCase):
|
|
def setUp(self):
|
|
self.server = TranscriptionServer()
|
|
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
|
|
self.server.cache_path = "~/.cache/whisper-live/"
|
|
self.server.backend = BackendType.FASTER_WHISPER
|
|
|
|
@mock.patch("websockets.WebSocketCommonProtocol")
|
|
def test_invalid_json_returns_false(self, mock_ws):
|
|
mock_ws.recv.return_value = "not valid json {{"
|
|
result = self.server.handle_new_connection(mock_ws, None, None, False)
|
|
self.assertFalse(result)
|
|
|
|
@mock.patch("websockets.WebSocketCommonProtocol")
|
|
def test_server_full_returns_false(self, mock_ws):
|
|
# Fill server
|
|
for i in range(4):
|
|
self.server.client_manager.add_client(MagicMock(), MagicMock())
|
|
|
|
mock_ws.recv.return_value = json.dumps({
|
|
"uid": "test",
|
|
"language": "en",
|
|
"task": "transcribe",
|
|
"model": "tiny.en",
|
|
})
|
|
result = self.server.handle_new_connection(mock_ws, None, None, False)
|
|
self.assertFalse(result)
|
|
|
|
|
|
class TestTranscriptionServerCleanup(unittest.TestCase):
|
|
def setUp(self):
|
|
self.server = TranscriptionServer()
|
|
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
|
|
|
|
def test_cleanup_removes_client(self):
|
|
ws = MagicMock()
|
|
client = MagicMock()
|
|
self.server.client_manager.add_client(ws, client)
|
|
self.cleanup_server = self.server
|
|
self.server.cleanup(ws)
|
|
self.assertNotIn(ws, self.server.client_manager.clients)
|
|
client.cleanup.assert_called_once()
|
|
|
|
|
|
class TestStreamTranscription(unittest.TestCase):
|
|
"""Tests for the SSE streaming endpoint (stream=true)."""
|
|
|
|
def _make_app(self):
|
|
"""Create a FastAPI app with the transcribe endpoint that has streaming support."""
|
|
from fastapi import FastAPI, UploadFile, Form
|
|
|
|
app = FastAPI()
|
|
server = TranscriptionServer()
|
|
|
|
@app.post("/v1/audio/transcriptions")
|
|
async def transcribe(
|
|
file: UploadFile,
|
|
stream: bool = Form(default=False),
|
|
language: str = Form(default=None),
|
|
response_format: str = Form(default="json"),
|
|
):
|
|
if stream:
|
|
return server._stream_transcription(
|
|
file, language, None, 0.0, None, None
|
|
)
|
|
return {"text": "non-streamed"}
|
|
|
|
return app
|
|
|
|
@patch("whisper_live.server.WhisperModel")
|
|
def test_stream_returns_sse_content_type(self, mock_model_cls):
|
|
mock_seg = MagicMock()
|
|
mock_seg.id = 0
|
|
mock_seg.start = 0.0
|
|
mock_seg.end = 1.0
|
|
mock_seg.text = " hello "
|
|
mock_seg.words = []
|
|
|
|
mock_info = MagicMock()
|
|
mock_info.language = "en"
|
|
mock_info.language_probability = 0.98
|
|
mock_info.duration = 1.0
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.transcribe.return_value = (iter([mock_seg]), mock_info)
|
|
mock_model_cls.return_value = mock_model
|
|
|
|
import io
|
|
from fastapi.testclient import TestClient
|
|
|
|
app = self._make_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/audio/transcriptions",
|
|
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
|
|
data={"stream": "true"},
|
|
)
|
|
self.assertEqual(resp.status_code, 200)
|
|
self.assertIn("text/event-stream", resp.headers.get("content-type", ""))
|
|
|
|
@patch("whisper_live.server.WhisperModel")
|
|
def test_stream_yields_segment_and_done(self, mock_model_cls):
|
|
mock_seg = MagicMock()
|
|
mock_seg.id = 0
|
|
mock_seg.start = 0.0
|
|
mock_seg.end = 1.5
|
|
mock_seg.text = " hello world "
|
|
mock_seg.words = []
|
|
|
|
mock_info = MagicMock()
|
|
mock_info.language = "en"
|
|
mock_info.language_probability = 0.95
|
|
mock_info.duration = 1.5
|
|
mock_model = MagicMock()
|
|
mock_model.transcribe.return_value = (iter([mock_seg]), mock_info)
|
|
mock_model_cls.return_value = mock_model
|
|
|
|
import io
|
|
from fastapi.testclient import TestClient
|
|
|
|
app = self._make_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/audio/transcriptions",
|
|
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
|
|
data={"stream": "true"},
|
|
)
|
|
body = resp.text
|
|
self.assertIn('"text": "hello world"', body)
|
|
self.assertIn("[DONE]", body)
|
|
|
|
@patch("whisper_live.server.WhisperModel")
|
|
def test_stream_multiple_segments(self, mock_model_cls):
|
|
segs = []
|
|
for i in range(3):
|
|
s = MagicMock()
|
|
s.id = i
|
|
s.start = float(i)
|
|
s.end = float(i + 1)
|
|
s.text = f" segment {i} "
|
|
s.words = []
|
|
segs.append(s)
|
|
|
|
mock_info = MagicMock()
|
|
mock_info.language = "en"
|
|
mock_info.language_probability = 0.99
|
|
mock_info.duration = 3.0
|
|
mock_model = MagicMock()
|
|
mock_model.transcribe.return_value = (iter(segs), mock_info)
|
|
mock_model_cls.return_value = mock_model
|
|
|
|
import io
|
|
from fastapi.testclient import TestClient
|
|
|
|
app = self._make_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/audio/transcriptions",
|
|
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
|
|
data={"stream": "true"},
|
|
)
|
|
body = resp.text
|
|
events = [line for line in body.split("\n") if line.startswith("data: ") and "[DONE]" not in line and '"type": "metadata"' not in line]
|
|
self.assertEqual(len(events), 3)
|
|
for i, event in enumerate(events):
|
|
data = json.loads(event.removeprefix("data: "))
|
|
self.assertEqual(data["text"], f"segment {i}")
|
|
|
|
@patch("whisper_live.server.WhisperModel", side_effect=RuntimeError("model error"))
|
|
def test_stream_error_yields_error_event(self, mock_model_cls):
|
|
import io
|
|
from fastapi.testclient import TestClient
|
|
|
|
app = self._make_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/audio/transcriptions",
|
|
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
|
|
data={"stream": "true"},
|
|
)
|
|
body = resp.text
|
|
self.assertIn('"error"', body)
|
|
self.assertIn("model error", body)
|
|
|
|
def test_non_stream_still_works(self):
|
|
import io
|
|
from fastapi.testclient import TestClient
|
|
|
|
app = self._make_app()
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/v1/audio/transcriptions",
|
|
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
|
|
data={"stream": "false"},
|
|
)
|
|
self.assertEqual(resp.status_code, 200)
|
|
self.assertEqual(resp.json()["text"], "non-streamed")
|
|
|
|
|
|
class TestRESTAPIParamWarnings(unittest.TestCase):
|
|
"""Test that unsupported OpenAI-compatible REST params produce warnings."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""Build a FastAPI test app by extracting the endpoint definition."""
|
|
import logging
|
|
from fastapi import FastAPI, UploadFile, Form, File
|
|
from fastapi.testclient import TestClient
|
|
from typing import Optional, List
|
|
|
|
app = FastAPI()
|
|
|
|
@app.post("/v1/audio/transcriptions")
|
|
async def transcribe(
|
|
file: UploadFile,
|
|
model: str = Form(default="whisper-1"),
|
|
language: Optional[str] = Form(default=None),
|
|
prompt: Optional[str] = Form(default=None),
|
|
response_format: str = Form(default="json"),
|
|
temperature: float = Form(default=0.0),
|
|
timestamp_granularities: Optional[List[str]] = Form(default=None),
|
|
chunking_strategy: Optional[str] = Form(default=None),
|
|
include: Optional[List[str]] = Form(default=None),
|
|
known_speaker_names: Optional[List[str]] = Form(default=None),
|
|
known_speaker_references: Optional[List[UploadFile]] = File(default=None),
|
|
stream: bool = Form(default=False),
|
|
):
|
|
ignored_params = []
|
|
if chunking_strategy:
|
|
ignored_params.append(f"chunking_strategy='{chunking_strategy}'")
|
|
if include:
|
|
ignored_params.append(f"include={include}")
|
|
if ignored_params:
|
|
logging.warning(f"Unsupported OpenAI params ignored: {', '.join(ignored_params)}")
|
|
# Return a JSON response with the ignored list for testing
|
|
return {"text": "test", "ignored": ignored_params}
|
|
|
|
cls.test_client = TestClient(app)
|
|
|
|
def _post(self, **extra_fields):
|
|
import io
|
|
data = {**extra_fields}
|
|
files = {"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")}
|
|
return self.test_client.post("/v1/audio/transcriptions", data=data, files=files)
|
|
|
|
def test_no_warnings_when_no_extra_params(self):
|
|
resp = self._post()
|
|
self.assertEqual(resp.status_code, 200)
|
|
self.assertEqual(resp.json()["ignored"], [])
|
|
|
|
def test_chunking_strategy_warning(self):
|
|
resp = self._post(chunking_strategy="auto")
|
|
self.assertEqual(resp.status_code, 200)
|
|
ignored = resp.json()["ignored"]
|
|
self.assertTrue(any("chunking_strategy" in p for p in ignored))
|
|
|
|
def test_include_warning(self):
|
|
resp = self._post(include="logprobs")
|
|
self.assertEqual(resp.status_code, 200)
|
|
ignored = resp.json()["ignored"]
|
|
self.assertTrue(any("include" in p for p in ignored))
|
|
|
|
def test_known_speaker_names_supported(self):
|
|
resp = self._post(known_speaker_names="alice")
|
|
self.assertEqual(resp.status_code, 200)
|
|
ignored = resp.json()["ignored"]
|
|
self.assertFalse(any("known_speaker_names" in p for p in ignored))
|
|
|
|
def test_multiple_ignored_params(self):
|
|
resp = self._post(chunking_strategy="auto", known_speaker_names="bob")
|
|
self.assertEqual(resp.status_code, 200)
|
|
ignored = resp.json()["ignored"]
|
|
self.assertEqual(len(ignored), 1)
|
|
|
|
|
|
class TestAPIKeyAuth(unittest.TestCase):
|
|
"""Test optional API key authentication middleware."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.testclient import TestClient
|
|
from fastapi.responses import JSONResponse as JSONR
|
|
|
|
app = FastAPI()
|
|
|
|
@app.middleware("http")
|
|
async def _check_api_key(request: Request, call_next):
|
|
auth = request.headers.get("Authorization", "")
|
|
if auth != "Bearer test-secret":
|
|
return JSONR({"error": "Invalid or missing API key"}, status_code=401)
|
|
return await call_next(request)
|
|
|
|
@app.get("/ping")
|
|
async def ping():
|
|
return {"status": "ok"}
|
|
|
|
cls.test_client = TestClient(app)
|
|
|
|
def test_missing_key_returns_401(self):
|
|
resp = self.test_client.get("/ping")
|
|
self.assertEqual(resp.status_code, 401)
|
|
|
|
def test_wrong_key_returns_401(self):
|
|
resp = self.test_client.get("/ping", headers={"Authorization": "Bearer wrong"})
|
|
self.assertEqual(resp.status_code, 401)
|
|
|
|
def test_correct_key_returns_200(self):
|
|
resp = self.test_client.get("/ping", headers={"Authorization": "Bearer test-secret"})
|
|
self.assertEqual(resp.status_code, 200)
|
|
self.assertEqual(resp.json()["status"], "ok")
|
|
|
|
|
|
class TestRateLimiting(unittest.TestCase):
|
|
"""Test per-IP rate limiting middleware."""
|
|
|
|
def _make_app(self, rpm_limit=3):
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.testclient import TestClient
|
|
from fastapi.responses import JSONResponse as JSONR
|
|
|
|
_rate_lock = threading.Lock()
|
|
_rate_buckets: dict = {}
|
|
|
|
app = FastAPI()
|
|
|
|
@app.middleware("http")
|
|
async def _rate_limit(request: Request, call_next):
|
|
client_ip = request.client.host if request.client else "unknown"
|
|
now = time.time()
|
|
with _rate_lock:
|
|
bucket = _rate_buckets.setdefault(client_ip, collections.deque())
|
|
while bucket and bucket[0] < now - 60:
|
|
bucket.popleft()
|
|
if len(bucket) >= rpm_limit:
|
|
return JSONR({"error": "Rate limit exceeded"}, status_code=429)
|
|
bucket.append(now)
|
|
return await call_next(request)
|
|
|
|
@app.get("/ping")
|
|
async def ping():
|
|
return {"status": "ok"}
|
|
|
|
return TestClient(app)
|
|
|
|
def test_within_limit_succeeds(self):
|
|
client = self._make_app(rpm_limit=3)
|
|
for _ in range(3):
|
|
resp = client.get("/ping")
|
|
self.assertEqual(resp.status_code, 200)
|
|
|
|
def test_exceeding_limit_returns_429(self):
|
|
client = self._make_app(rpm_limit=3)
|
|
for _ in range(3):
|
|
client.get("/ping")
|
|
resp = client.get("/ping")
|
|
self.assertEqual(resp.status_code, 429)
|
|
self.assertIn("Rate limit", resp.json()["error"])
|
|
|
|
|
|
class TestWebSocketAuth(unittest.TestCase):
|
|
"""Tests for the WebSocket process_request auth callback."""
|
|
|
|
def _make_auth_handler(self, api_key):
|
|
"""Build the same auth function the server creates."""
|
|
def _ws_auth(path, request_headers):
|
|
auth = request_headers.get("Authorization", "")
|
|
token_param = None
|
|
if "?" in path:
|
|
from urllib.parse import urlparse, parse_qs
|
|
parsed = urlparse(path)
|
|
token_param = parse_qs(parsed.query).get("token", [None])[0]
|
|
if auth == f"Bearer {api_key}" or token_param == api_key:
|
|
return None
|
|
return (401, [("Content-Type", "text/plain")], b"Unauthorized\n")
|
|
return _ws_auth
|
|
|
|
def test_valid_bearer_token(self):
|
|
handler = self._make_auth_handler("my-secret")
|
|
result = handler("/", {"Authorization": "Bearer my-secret"})
|
|
self.assertIsNone(result)
|
|
|
|
def test_invalid_bearer_token(self):
|
|
handler = self._make_auth_handler("my-secret")
|
|
result = handler("/", {"Authorization": "Bearer wrong"})
|
|
self.assertEqual(result[0], 401)
|
|
|
|
def test_missing_auth_header(self):
|
|
handler = self._make_auth_handler("my-secret")
|
|
result = handler("/", {})
|
|
self.assertEqual(result[0], 401)
|
|
|
|
def test_valid_query_token(self):
|
|
handler = self._make_auth_handler("my-secret")
|
|
result = handler("/?token=my-secret", {})
|
|
self.assertIsNone(result)
|
|
|
|
def test_invalid_query_token(self):
|
|
handler = self._make_auth_handler("my-secret")
|
|
result = handler("/?token=wrong", {})
|
|
self.assertEqual(result[0], 401)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|