import io import json import queue import sys import tempfile import unittest import wave from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch import trainer_server as trainer def silent_wav_bytes(duration_s: float = 0.25) -> bytes: output = io.BytesIO() with wave.open(output, "wb") as wav_file: wav_file.setnchannels(1) wav_file.setsampwidth(2) wav_file.setframerate(16000) wav_file.writeframes(b"\x00\x00" * int(16000 * duration_s)) return output.getvalue() class AutoTrainTests(unittest.TestCase): def clear_review_queue(self): while True: try: trainer.AUTO_TRAIN_REVIEW_QUEUE.get_nowait() except queue.Empty: break else: trainer.AUTO_TRAIN_REVIEW_QUEUE.task_done() trainer.AUTO_TRAIN_QUEUED_FILES.clear() def setUp(self): self.clear_review_queue() self.tempdir = tempfile.TemporaryDirectory() root = Path(self.tempdir.name) self.original_paths = ( trainer.CAPTURED_DIR, trainer.NEGATIVE_DIR, trainer.PERSONAL_DIR, trainer.AUTO_TRAIN_CONFIG_FILE, trainer.AUTO_TRAIN_STATE_FILE, ) trainer.CAPTURED_DIR = root / "captured_audio" trainer.NEGATIVE_DIR = root / "negative_samples" trainer.PERSONAL_DIR = root / "personal_samples" trainer.AUTO_TRAIN_CONFIG_FILE = root / "auto_train_config.json" trainer.AUTO_TRAIN_STATE_FILE = root / "auto_train_state.json" for directory in (trainer.CAPTURED_DIR, trainer.NEGATIVE_DIR, trainer.PERSONAL_DIR): directory.mkdir(parents=True) self.original_config = dict(trainer.AUTO_TRAIN_CONFIG) self.original_state = dict(trainer.AUTO_TRAIN_STATE) trainer.AUTO_TRAIN_CONFIG.clear() trainer.AUTO_TRAIN_CONFIG.update( trainer._normalize_auto_train_config( { "enabled": True, "wake_phrase": "hey tater", "language": "en", "tater_url": "http://127.0.0.1:8501", } ) ) trainer.AUTO_TRAIN_STATE.clear() trainer.AUTO_TRAIN_STATE.update(trainer.AUTO_TRAIN_DEFAULT_STATE) def tearDown(self): ( trainer.CAPTURED_DIR, trainer.NEGATIVE_DIR, trainer.PERSONAL_DIR, trainer.AUTO_TRAIN_CONFIG_FILE, trainer.AUTO_TRAIN_STATE_FILE, ) = self.original_paths trainer.AUTO_TRAIN_CONFIG.clear() trainer.AUTO_TRAIN_CONFIG.update(self.original_config) trainer.AUTO_TRAIN_STATE.clear() trainer.AUTO_TRAIN_STATE.update(self.original_state) self.clear_review_queue() self.tempdir.cleanup() def add_capture( self, name: str = "wake.wav", wake_word: str = "hey_tater", event_type: str = "wake_detected", blocked_by_vad: bool = False, ) -> Path: audio_path = trainer.CAPTURED_DIR / name audio_path.write_bytes(silent_wav_bytes()) trainer._write_sidecar_json( audio_path, { "original_name": name, "wake_word": wake_word, "event_type": event_type, "blocked_by_vad": blocked_by_vad, "review_status": "pending", }, ) return audio_path def test_phrase_matching_normalizes_case_punctuation_and_underscores(self): self.assertTrue(trainer._transcript_contains_wake_phrase("Okay, HEY TATER!", "hey_tater")) self.assertFalse(trainer._transcript_contains_wake_phrase("Turn on the television", "hey tater")) def test_stt_engine_selection_uses_managed_models(self): config = trainer._normalize_auto_train_config( { "stt_engine": "parakeet-onnx", "stt_model": "user/should-not-be-used", "stt_device": "cpu", "stt_compute_type": "float32", } ) self.assertEqual(config["stt_engine"], trainer.STT_ENGINE_PARAKEET_ONNX) self.assertNotIn("stt_model", config) self.assertNotIn("stt_device", config) self.assertNotIn("stt_compute_type", config) self.assertEqual( trainer._managed_stt_model(config["stt_engine"], "en"), trainer.DEFAULT_PARAKEET_ONNX_MODEL, ) self.assertEqual( trainer._managed_stt_model(trainer.STT_ENGINE_FASTER_WHISPER, "de"), trainer.DEFAULT_FASTER_WHISPER_MULTILINGUAL_MODEL, ) def test_stt_router_supports_both_nvidia_engines(self): audio_path = Path("wake.wav") with ( patch.object(trainer, "_transcribe_capture_with_faster_whisper", return_value="faster") as faster, patch.object(trainer, "_transcribe_capture_with_parakeet", return_value="parakeet") as parakeet, ): self.assertEqual( trainer._transcribe_capture( audio_path, engine=trainer.STT_ENGINE_FASTER_WHISPER, language="en", ), "faster", ) self.assertEqual( trainer._transcribe_capture( audio_path, engine=trainer.STT_ENGINE_PARAKEET_ONNX, language="en", ), "parakeet", ) faster.assert_called_once() parakeet.assert_called_once() def test_parakeet_loader_prefers_cuda_then_cpu(self): fake_model = object() fake_onnx_asr = SimpleNamespace(load_model=Mock(return_value=fake_model)) with ( patch.dict(sys.modules, {"onnx_asr": fake_onnx_asr}), patch.object( trainer, "_parakeet_onnx_providers", return_value=["CUDAExecutionProvider", "CPUExecutionProvider"], ), ): with trainer.PARAKEET_ONNX_MODEL_LOCK: trainer.PARAKEET_ONNX_MODEL_CACHE.clear() loaded = trainer._load_parakeet_onnx_model() self.assertIs(loaded, fake_model) fake_onnx_asr.load_model.assert_called_once_with( trainer.DEFAULT_PARAKEET_ONNX_MODEL, str(trainer.AUTO_TRAIN_MODEL_DIR), quantization="int8", providers=["CUDAExecutionProvider", "CPUExecutionProvider"], ) def test_ui_exposes_engine_selector_without_manual_runtime_fields(self): source = (Path(__file__).resolve().parents[1] / "static" / "index.html").read_text( encoding="utf-8" ) self.assertIn('id="autoSttEngine"', source) self.assertNotIn('id="autoSttModel"', source) self.assertNotIn('id="autoSttDevice"', source) self.assertNotIn('id="autoSttComputeType"', source) def test_phrase_miss_moves_wake_trigger_to_negative_samples(self): self.add_capture() with patch.object(trainer, "_transcribe_capture", return_value="turn on the kitchen lights"): trainer._auto_review_capture("wake.wav") self.assertFalse((trainer.CAPTURED_DIR / "wake.wav").exists()) negatives = list(trainer.NEGATIVE_DIR.glob("*.wav")) self.assertEqual(len(negatives), 1) metadata = trainer._load_sidecar_json(negatives[0]) self.assertTrue(metadata["auto_negative"]) self.assertEqual(metadata["review_status"], "auto_approved_negative") self.assertEqual(metadata["transcript"], "turn on the kitchen lights") self.assertEqual(metadata["auto_review_stt_engine"], "faster_whisper") self.assertEqual(metadata["auto_review_stt_model"], "small.en") self.assertEqual(trainer.AUTO_TRAIN_STATE["pending_negative_count"], 1) def test_matching_phrase_stays_in_manual_review_inbox(self): audio_path = self.add_capture() with patch.object(trainer, "_transcribe_capture", return_value="hey tater turn on the lights"): trainer._auto_review_capture("wake.wav") self.assertTrue(audio_path.exists()) self.assertFalse(list(trainer.NEGATIVE_DIR.glob("*.wav"))) metadata = trainer._load_sidecar_json(audio_path) self.assertEqual(metadata["auto_review_status"], "wake_phrase_detected") self.assertEqual(trainer.AUTO_TRAIN_STATE["pending_negative_count"], 0) def test_matching_phrase_is_deleted_when_cleanup_is_enabled(self): audio_path = self.add_capture() trainer.AUTO_TRAIN_CONFIG["delete_confirmed_wakes"] = True with patch.object( trainer, "_transcribe_capture", return_value="hey tater turn on the lights", ): trainer._auto_review_capture("wake.wav") self.assertFalse(audio_path.exists()) self.assertFalse(audio_path.with_suffix(".json").exists()) self.assertFalse(list(trainer.PERSONAL_DIR.glob("*.wav"))) self.assertFalse(list(trainer.NEGATIVE_DIR.glob("*.wav"))) self.assertEqual(trainer.AUTO_TRAIN_STATE["last_review_result"], "deleted_confirmed_wake") def test_cleanup_processes_previously_confirmed_wake_without_retranscribing(self): audio_path = self.add_capture() metadata = trainer._load_sidecar_json(audio_path) metadata.update( { "auto_review_status": "wake_phrase_detected", "transcript": "hey tater", } ) trainer._write_sidecar_json(audio_path, metadata) trainer.AUTO_TRAIN_CONFIG["delete_confirmed_wakes"] = True self.assertEqual(trainer._queue_pending_auto_reviews(), 1) with patch.object(trainer, "_transcribe_capture") as transcribe: trainer._auto_review_capture("wake.wav") transcribe.assert_not_called() self.assertFalse(audio_path.exists()) self.assertEqual(trainer.AUTO_TRAIN_STATE["last_review_transcript"], "hey tater") def test_close_miss_is_not_transcribed_by_default(self): audio_path = self.add_capture(event_type="close_miss") with patch.object(trainer, "_transcribe_capture") as transcribe: trainer._auto_review_capture("wake.wav") transcribe.assert_not_called() self.assertTrue(audio_path.exists()) self.assertFalse(trainer._load_sidecar_json(audio_path).get("auto_review_status")) def test_existing_close_miss_is_queued_when_promotion_is_enabled(self): self.add_capture(event_type="close_miss") self.assertEqual(trainer._queue_pending_auto_reviews(), 0) trainer.AUTO_TRAIN_CONFIG["promote_close_misses"] = True self.assertEqual(trainer._queue_pending_auto_reviews(), 1) def test_close_miss_with_phrase_is_promoted_when_enabled(self): self.add_capture(event_type="close_miss") trainer.AUTO_TRAIN_CONFIG["promote_close_misses"] = True with patch.object(trainer, "_transcribe_capture", return_value="hey tater"): trainer._auto_review_capture("wake.wav") self.assertFalse((trainer.CAPTURED_DIR / "wake.wav").exists()) positives = list(trainer.PERSONAL_DIR.glob("*.wav")) self.assertEqual(len(positives), 1) metadata = trainer._load_sidecar_json(positives[0]) self.assertTrue(metadata["auto_positive"]) self.assertEqual(metadata["review_status"], "auto_approved_personal") self.assertEqual(metadata["transcript"], "hey tater") self.assertFalse(list(trainer.NEGATIVE_DIR.glob("*.wav"))) self.assertEqual(trainer.AUTO_TRAIN_STATE["pending_negative_count"], 0) def test_close_miss_without_phrase_stays_in_inbox(self): audio_path = self.add_capture(event_type="close_miss") trainer.AUTO_TRAIN_CONFIG["promote_close_misses"] = True with patch.object( trainer, "_transcribe_capture", return_value="turn on the lights", ): trainer._auto_review_capture("wake.wav") self.assertTrue(audio_path.exists()) self.assertFalse(list(trainer.PERSONAL_DIR.glob("*.wav"))) self.assertFalse(list(trainer.NEGATIVE_DIR.glob("*.wav"))) metadata = trainer._load_sidecar_json(audio_path) self.assertEqual(metadata["auto_review_status"], "close_miss_phrase_not_detected") def test_vad_blocked_close_miss_is_never_transcribed(self): audio_path = self.add_capture(event_type="close_miss", blocked_by_vad=True) trainer.AUTO_TRAIN_CONFIG["promote_close_misses"] = True with patch.object(trainer, "_transcribe_capture") as transcribe: trainer._auto_review_capture("wake.wav") transcribe.assert_not_called() self.assertTrue(audio_path.exists()) self.assertFalse(trainer._load_sidecar_json(audio_path).get("auto_review_status")) def test_capture_for_another_wake_word_is_not_transcribed(self): audio_path = self.add_capture(wake_word="computer") with patch.object(trainer, "_transcribe_capture") as transcribe: trainer._auto_review_capture("wake.wav") transcribe.assert_not_called() self.assertTrue(audio_path.exists()) metadata = trainer._load_sidecar_json(audio_path) self.assertEqual(metadata["auto_review_status"], "different_wake_phrase") def test_due_schedule_starts_training_after_minimum_negatives(self): trainer.AUTO_TRAIN_CONFIG["schedule_hours"] = 24 trainer.AUTO_TRAIN_CONFIG["minimum_new_negatives"] = 3 trainer.AUTO_TRAIN_STATE["pending_negative_count"] = 3 trainer.AUTO_TRAIN_STATE["next_run_at"] = "2000-01-01T00:00:00+00:00" with patch.object(trainer, "_start_auto_training", return_value={"ok": True, "started": True}) as start: trainer._maybe_run_scheduled_auto_training() start.assert_called_once_with() self.assertTrue(trainer.AUTO_TRAIN_STATE["next_run_at"]) def test_tater_notification_sets_new_word_globally_with_token(self): trainer.AUTO_TRAIN_CONFIG.update( { "notify_satellites": True, "tater_url": "http://127.0.0.1:8501", "tater_link_token": "secret-token", } ) class Response: def __enter__(self): return self def __exit__(self, *_args): return False def read(self): return b'{"push":{"count":4}}' trained_word = { "key": "hey_tater", "wake_word": "Hey Tater", "json_url": "http://10.4.20.210:8789/api/trained_wake_words/hey_tater.json", } with ( patch.object(trainer, "_advertised_base_url", return_value="http://10.4.20.210:8789"), patch.object(trainer, "_list_trained_wake_words", return_value=[trained_word]) as catalog, patch.object(trainer, "urlopen", return_value=Response()) as open_url, ): result = trainer._notify_tater_satellites("hey_tater") self.assertTrue(result["ok"]) self.assertEqual(result["count"], 4) self.assertEqual(result["wake_word"], "Hey Tater") self.assertEqual(result["wake_word_url"], trained_word["json_url"]) catalog.assert_called_once_with("http://10.4.20.210:8789") self.assertEqual(open_url.call_count, 1) request = open_url.call_args.args[0] self.assertEqual(request.full_url, "http://127.0.0.1:8501/api/tater/satellite/v1/trainer/wake-word") self.assertEqual(request.get_method(), "POST") self.assertEqual(request.get_header("X-tater-trainer-token"), "secret-token") self.assertEqual( json.loads(request.data), { "wake_word_name": "hey_tater", "wake_word_url": trained_word["json_url"], }, ) def test_tater_notification_fails_when_trained_word_is_missing(self): trainer.AUTO_TRAIN_CONFIG["tater_link_token"] = "secret-token" with ( patch.object(trainer, "_advertised_base_url", return_value="http://10.4.20.210:8789"), patch.object(trainer, "_list_trained_wake_words", return_value=[]), patch.object(trainer, "urlopen") as open_url, ): result = trainer._notify_tater_satellites("missing_word") self.assertFalse(result["ok"]) self.assertIn("missing_word", result["error"]) open_url.assert_not_called() def test_tater_notification_requires_secure_link(self): trainer.AUTO_TRAIN_CONFIG["tater_link_token"] = "" with patch.object(trainer, "urlopen") as open_url: result = trainer._notify_tater_satellites("hey_tater") self.assertFalse(result["ok"]) self.assertIn("not linked", result["error"]) open_url.assert_not_called() def test_claim_tater_link_uses_tater_code_and_keeps_token_private(self): class Response: def __enter__(self): return self def __exit__(self, *_args): return False def read(self, *_args): return json.dumps( { "ok": True, "token": "a" * 43, "tater_name": "Tater", "linked_at": "2026-07-24T12:00:00+00:00", } ).encode("utf-8") with ( patch.object(trainer, "_advertised_base_url", return_value="http://10.4.20.210:8789"), patch.object(trainer, "urlopen", return_value=Response()) as open_url, ): result = trainer._claim_tater_link("http://127.0.0.1:8501", "ABCD-EFGH") self.assertTrue(result["linked"]) self.assertEqual(trainer.AUTO_TRAIN_CONFIG["tater_link_token"], "a" * 43) self.assertNotIn("tater_link_token", trainer._public_auto_train_config()) request = open_url.call_args.args[0] self.assertEqual( request.full_url, "http://127.0.0.1:8501/api/tater/satellite/v1/trainer/link/claim", ) payload = json.loads(request.data) self.assertEqual(payload["pairing_code"], "ABCDEFGH") self.assertEqual(payload["publish_base_url"], "http://10.4.20.210:8789") self.assertTrue(payload["trainer_id"]) def test_advertised_url_uses_non_loopback_browser_host(self): request = SimpleNamespace( base_url="http://192.168.1.50:8789/", url=SimpleNamespace(hostname="192.168.1.50", scheme="http", port=8789), ) self.assertEqual(trainer._advertised_base_url(request), "http://192.168.1.50:8789") def test_advertised_url_replaces_localhost_with_discovered_lan_host(self): request = SimpleNamespace( base_url="http://127.0.0.1:8789/", url=SimpleNamespace(hostname="127.0.0.1", scheme="http", port=8789), ) with patch.object(trainer, "_discover_lan_ipv4", return_value="192.168.1.60"): self.assertEqual(trainer._advertised_base_url(request), "http://192.168.1.60:8789") def test_configured_public_url_takes_precedence(self): trainer.AUTO_TRAIN_CONFIG["advertised_base_url"] = "http://trainer.local:8789" request = SimpleNamespace( base_url="http://127.0.0.1:8789/", url=SimpleNamespace(hostname="127.0.0.1", scheme="http", port=8789), ) self.assertEqual(trainer._advertised_base_url(request), "http://trainer.local:8789") def test_faster_whisper_auto_runtime_prefers_cuda_and_float16(self): fake_ctranslate2 = SimpleNamespace(get_cuda_device_count=lambda: 1) with patch.dict(sys.modules, {"ctranslate2": fake_ctranslate2}): self.assertEqual( trainer._resolve_faster_whisper_runtime("auto", "auto"), ("cuda", "float16"), ) def test_faster_whisper_auto_runtime_falls_back_to_cpu_int8(self): fake_ctranslate2 = SimpleNamespace(get_cuda_device_count=lambda: 0) with patch.dict(sys.modules, {"ctranslate2": fake_ctranslate2}): self.assertEqual( trainer._resolve_faster_whisper_runtime("auto", "auto"), ("cpu", "int8"), ) def test_faster_whisper_transcription_joins_segments_and_records_runtime(self): fake_model = SimpleNamespace() fake_model.transcribe = Mock( return_value=( iter([SimpleNamespace(text=" turn on "), SimpleNamespace(text="the lights ")]), SimpleNamespace(), ) ) with ( patch.object(trainer, "_resolve_faster_whisper_runtime", return_value=("cuda", "float16")), patch.object(trainer, "_load_faster_whisper_model", return_value=fake_model), ): transcript = trainer._transcribe_capture_with_faster_whisper( Path("wake.wav"), model="small.en", language="en", ) self.assertEqual(transcript, "turn on the lights") fake_model.transcribe.assert_called_once_with( "wake.wav", language="en", beam_size=1, condition_on_previous_text=False, ) self.assertEqual(trainer.AUTO_TRAIN_STATE["last_stt_device"], "cuda") self.assertEqual(trainer.AUTO_TRAIN_STATE["last_stt_compute_type"], "float16") def test_train_status_reads_and_increments_training_log_tail(self): log_path = Path(self.tempdir.name) / "training.log" log_path.write_text("first\nsecond\nthird\n", encoding="utf-8") with trainer.STATE_LOCK: original_training = dict(trainer.STATE["training"]) trainer.STATE["training"].update( { "log_path": str(log_path), "last_sent_tail": [], "last_log_size": 0, } ) try: with ( patch.object(trainer, "TRAIN_LOG_TAIL_LINES", 2), patch.object(trainer, "TRAIN_LOG_MAX_BYTES", 1024), ): first_status = trainer.train_status() self.assertEqual(first_status["training"]["log_lines"], ["second", "third"]) self.assertEqual(first_status["training"]["log_text"], "second\nthird") with log_path.open("a", encoding="utf-8") as log_file: log_file.write("fourth\n") next_status = trainer.train_status() self.assertEqual(next_status["training"]["log_lines"], ["third", "fourth"]) self.assertEqual(next_status["training"]["log_text"], "fourth") finally: with trainer.STATE_LOCK: trainer.STATE["training"].clear() trainer.STATE["training"].update(original_training) if __name__ == "__main__": unittest.main()