135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
import base64
|
|
import threading
|
|
import unittest
|
|
from dataclasses import FrozenInstanceError
|
|
|
|
import numpy as np
|
|
|
|
from blurt_client import (
|
|
OPENAI_TRANSCRIPTIONS_URL,
|
|
OPENROUTER_TRANSCRIPTIONS_URL,
|
|
TranscriptionClient,
|
|
TranscriptionConfig,
|
|
TranscriptionError,
|
|
transcription_timeout,
|
|
)
|
|
|
|
|
|
class FakeResponse:
|
|
status_code = 200
|
|
headers = {"content-type": "application/json"}
|
|
content = b"{}"
|
|
text = "{}"
|
|
ok = True
|
|
|
|
def __init__(self, result=None):
|
|
self.result = {"text": " transcript"} if result is None else result
|
|
|
|
def raise_for_status(self):
|
|
return None
|
|
|
|
def json(self):
|
|
return self.result
|
|
|
|
|
|
class FakeSession:
|
|
def __init__(self, result=None):
|
|
self.calls = []
|
|
self.response = FakeResponse(result)
|
|
self.closed = False
|
|
|
|
def post(self, url, **kwargs):
|
|
self.calls.append((url, kwargs))
|
|
return self.response
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
class ClientTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.sessions = []
|
|
|
|
def factory():
|
|
session = FakeSession()
|
|
self.sessions.append(session)
|
|
return session
|
|
|
|
self.client = TranscriptionClient(factory)
|
|
self.audio = np.zeros(32_000, dtype=np.int16)
|
|
|
|
def tearDown(self):
|
|
self.client.close()
|
|
|
|
def test_config_is_frozen_and_hides_key(self):
|
|
config = TranscriptionConfig("model", "en", "openai", "secret")
|
|
self.assertNotIn("secret", repr(config))
|
|
with self.assertRaises(FrozenInstanceError):
|
|
config.model = "changed"
|
|
|
|
def test_openai_multipart_payload_and_diarization(self):
|
|
config = TranscriptionConfig("gpt-4o-transcribe-diarize", "de", "openai", "key")
|
|
self.assertEqual(self.client.transcribe(self.audio, config), "transcript")
|
|
url, call = self.sessions[0].calls[0]
|
|
self.assertEqual(url, OPENAI_TRANSCRIPTIONS_URL)
|
|
self.assertEqual(call["headers"], {"Authorization": "Bearer key"})
|
|
self.assertEqual(call["files"]["file"][0::2], ("blurt.wav", "audio/wav"))
|
|
self.assertTrue(call["files"]["file"][1].startswith(b"RIFF"))
|
|
self.assertEqual(call["data"], {
|
|
"model": config.model, "language": "de",
|
|
"response_format": "diarized_json", "chunking_strategy": "auto",
|
|
})
|
|
self.assertEqual(call["timeout"], (10.0, 31.0))
|
|
|
|
def test_openrouter_base64_payload_and_headers(self):
|
|
config = TranscriptionConfig("openai/whisper-large-v3", "en", "openrouter", "router-key")
|
|
self.client.transcribe(self.audio, config)
|
|
url, call = self.sessions[0].calls[0]
|
|
self.assertEqual(url, OPENROUTER_TRANSCRIPTIONS_URL)
|
|
self.assertEqual(call["headers"]["Authorization"], "Bearer router-key")
|
|
self.assertEqual(call["headers"]["X-OpenRouter-Title"], "Blurt")
|
|
payload = call["json"]
|
|
self.assertEqual(payload["model"], config.model)
|
|
self.assertEqual(payload["language"], "en")
|
|
self.assertEqual(payload["input_audio"]["format"], "wav")
|
|
self.assertTrue(base64.b64decode(payload["input_audio"]["data"]).startswith(b"RIFF"))
|
|
|
|
def test_timeout_scales_and_caps(self):
|
|
self.assertEqual(transcription_timeout(0), (10.0, 30.0))
|
|
self.assertEqual(transcription_timeout(60), (10.0, 60.0))
|
|
self.assertEqual(transcription_timeout(10_000), (10.0, 300.0))
|
|
|
|
def test_thread_local_sessions_are_reused_and_closed(self):
|
|
config = TranscriptionConfig("whisper-1", "en", "openai", "key")
|
|
self.client.transcribe(self.audio, config)
|
|
self.client.transcribe(self.audio, config)
|
|
worker = threading.Thread(target=self.client.transcribe, args=(self.audio, config))
|
|
worker.start()
|
|
worker.join()
|
|
self.assertEqual(len(self.sessions), 2)
|
|
self.assertEqual(len(self.sessions[0].calls), 2)
|
|
self.client.close()
|
|
self.assertTrue(all(session.closed for session in self.sessions))
|
|
with self.assertRaises(RuntimeError):
|
|
self.client.transcribe(self.audio, config)
|
|
|
|
def test_empty_response_raises_provider_message(self):
|
|
session = FakeSession({"error": {"message": "bad audio"}})
|
|
client = TranscriptionClient(lambda: session)
|
|
with self.assertRaisesRegex(TranscriptionError, "bad audio"):
|
|
client.transcribe(
|
|
self.audio, TranscriptionConfig("whisper-1", "en", "openai", "key")
|
|
)
|
|
client.close()
|
|
|
|
def test_rejects_wrong_audio_shape_or_dtype(self):
|
|
config = TranscriptionConfig("whisper-1", "en", "openai", "key")
|
|
with self.assertRaises(ValueError):
|
|
self.client.transcribe(np.zeros(10, dtype=np.float32), config)
|
|
with self.assertRaises(ValueError):
|
|
self.client.transcribe(np.zeros((10, 1), dtype=np.int16), config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|