457 lines
16 KiB
Python
457 lines
16 KiB
Python
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
SRC = ROOT / "src"
|
|
if str(SRC) not in sys.path:
|
|
sys.path.insert(0, str(SRC))
|
|
|
|
import aman
|
|
from config import Config, VocabularyReplacement
|
|
|
|
|
|
class FakeDesktop:
|
|
def __init__(self):
|
|
self.inject_calls = []
|
|
self.quit_calls = 0
|
|
self.clipboard_text = ""
|
|
self.popup_open = False
|
|
self.popup_text = ""
|
|
self.popup_statuses = []
|
|
self.popup_callbacks = {}
|
|
self.popup_close_calls = 0
|
|
self.focus_restore_calls = 0
|
|
self.cancel_listener_active = False
|
|
self.cancel_listener_callback = None
|
|
|
|
def inject_text(
|
|
self,
|
|
text: str,
|
|
backend: str,
|
|
*,
|
|
remove_transcription_from_clipboard: bool = False,
|
|
) -> None:
|
|
self.inject_calls.append((text, backend, remove_transcription_from_clipboard))
|
|
|
|
def read_clipboard_text(self) -> str | None:
|
|
return self.clipboard_text
|
|
|
|
def write_clipboard_text(self, text: str) -> None:
|
|
self.clipboard_text = text
|
|
|
|
def open_edit_popup(
|
|
self,
|
|
initial_text: str,
|
|
*,
|
|
on_submit,
|
|
on_copy,
|
|
on_cancel,
|
|
) -> None:
|
|
self.popup_open = True
|
|
self.popup_text = initial_text
|
|
self.popup_callbacks = {
|
|
"submit": on_submit,
|
|
"copy": on_copy,
|
|
"cancel": on_cancel,
|
|
}
|
|
|
|
def close_edit_popup(self) -> None:
|
|
self.popup_open = False
|
|
self.popup_close_calls += 1
|
|
|
|
def get_edit_popup_text(self) -> str:
|
|
return self.popup_text
|
|
|
|
def set_edit_popup_text(self, text: str) -> None:
|
|
self.popup_text = text
|
|
|
|
def set_edit_popup_status(self, status: str) -> None:
|
|
self.popup_statuses.append(status)
|
|
|
|
def restore_previous_focus(self) -> bool:
|
|
self.focus_restore_calls += 1
|
|
return True
|
|
|
|
def start_cancel_listener(self, callback) -> None:
|
|
self.cancel_listener_active = True
|
|
self.cancel_listener_callback = callback
|
|
|
|
def stop_cancel_listener(self) -> None:
|
|
self.cancel_listener_active = False
|
|
self.cancel_listener_callback = None
|
|
|
|
def request_quit(self) -> None:
|
|
self.quit_calls += 1
|
|
|
|
|
|
class FakeSegment:
|
|
def __init__(self, text: str):
|
|
self.text = text
|
|
|
|
|
|
class FakeModel:
|
|
def __init__(self, text: str = "hello world"):
|
|
self.text = text
|
|
self.last_kwargs = {}
|
|
|
|
def transcribe(self, _audio, language=None, vad_filter=None):
|
|
self.last_kwargs = {
|
|
"language": language,
|
|
"vad_filter": vad_filter,
|
|
}
|
|
return [FakeSegment(self.text)], self.last_kwargs
|
|
|
|
|
|
class FakeHintModel:
|
|
def __init__(self, text: str = "hello world"):
|
|
self.text = text
|
|
self.last_kwargs = {}
|
|
|
|
def transcribe(
|
|
self,
|
|
_audio,
|
|
language=None,
|
|
vad_filter=None,
|
|
hotwords=None,
|
|
initial_prompt=None,
|
|
):
|
|
self.last_kwargs = {
|
|
"language": language,
|
|
"vad_filter": vad_filter,
|
|
"hotwords": hotwords,
|
|
"initial_prompt": initial_prompt,
|
|
}
|
|
return [FakeSegment(self.text)], self.last_kwargs
|
|
|
|
|
|
class FakeAIProcessor:
|
|
def __init__(self):
|
|
self.edit_calls = []
|
|
|
|
def process(self, text, lang="en", **_kwargs):
|
|
return text
|
|
|
|
def process_edit(
|
|
self,
|
|
current_text,
|
|
latest_instruction,
|
|
instruction_history,
|
|
lang="en",
|
|
**_kwargs,
|
|
):
|
|
self.edit_calls.append(
|
|
{
|
|
"current_text": current_text,
|
|
"latest_instruction": latest_instruction,
|
|
"instruction_history": list(instruction_history),
|
|
"lang": lang,
|
|
}
|
|
)
|
|
return f"{current_text} [{latest_instruction}]"
|
|
|
|
|
|
class FakeAudio:
|
|
def __init__(self, size: int):
|
|
self.size = size
|
|
|
|
|
|
class DaemonTests(unittest.TestCase):
|
|
def _config(self) -> Config:
|
|
cfg = Config()
|
|
return cfg
|
|
|
|
def _build_daemon(
|
|
self,
|
|
desktop: FakeDesktop,
|
|
model: FakeModel | FakeHintModel,
|
|
*,
|
|
cfg: Config | None = None,
|
|
verbose: bool = False,
|
|
) -> aman.Daemon:
|
|
active_cfg = cfg if cfg is not None else self._config()
|
|
with patch("aman._build_whisper_model", return_value=model), patch(
|
|
"aman.LlamaProcessor", return_value=FakeAIProcessor()
|
|
):
|
|
return aman.Daemon(active_cfg, desktop, verbose=verbose)
|
|
|
|
def _wait_until(self, predicate, timeout: float = 1.0):
|
|
end = time.time() + timeout
|
|
while time.time() < end:
|
|
if predicate():
|
|
return True
|
|
time.sleep(0.01)
|
|
return predicate()
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_toggle_start_stop_injects_text(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
daemon = self._build_daemon(desktop, FakeModel(), verbose=False)
|
|
daemon._start_stop_worker = (
|
|
lambda stream, record, trigger, process_audio: daemon._stop_and_process(
|
|
stream, record, trigger, process_audio
|
|
)
|
|
)
|
|
|
|
daemon.toggle()
|
|
self.assertEqual(daemon.get_state(), aman.State.RECORDING)
|
|
|
|
daemon.toggle()
|
|
|
|
self.assertEqual(daemon.get_state(), aman.State.IDLE)
|
|
self.assertEqual(desktop.inject_calls, [("hello world", "clipboard", False)])
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_shutdown_stops_recording_without_injection(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
daemon = self._build_daemon(desktop, FakeModel(), verbose=False)
|
|
daemon._start_stop_worker = (
|
|
lambda stream, record, trigger, process_audio: daemon._stop_and_process(
|
|
stream, record, trigger, process_audio
|
|
)
|
|
)
|
|
|
|
daemon.toggle()
|
|
self.assertEqual(daemon.get_state(), aman.State.RECORDING)
|
|
|
|
self.assertTrue(daemon.shutdown(timeout=0.2))
|
|
self.assertEqual(daemon.get_state(), aman.State.IDLE)
|
|
self.assertEqual(desktop.inject_calls, [])
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_dictionary_replacement_applies_after_ai(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
model = FakeModel(text="good morning martha")
|
|
cfg = self._config()
|
|
cfg.vocabulary.replacements = [VocabularyReplacement(source="Martha", target="Marta")]
|
|
|
|
daemon = self._build_daemon(desktop, model, cfg=cfg, verbose=False)
|
|
daemon._start_stop_worker = (
|
|
lambda stream, record, trigger, process_audio: daemon._stop_and_process(
|
|
stream, record, trigger, process_audio
|
|
)
|
|
)
|
|
|
|
daemon.toggle()
|
|
daemon.toggle()
|
|
|
|
self.assertEqual(desktop.inject_calls, [("good morning Marta", "clipboard", False)])
|
|
|
|
def test_transcribe_skips_hints_when_model_does_not_support_them(self):
|
|
desktop = FakeDesktop()
|
|
model = FakeModel(text="hello")
|
|
cfg = self._config()
|
|
cfg.vocabulary.terms = ["Docker", "Systemd"]
|
|
|
|
daemon = self._build_daemon(desktop, model, cfg=cfg, verbose=False)
|
|
|
|
result = daemon._transcribe(object())
|
|
|
|
self.assertEqual(result, "hello")
|
|
self.assertNotIn("hotwords", model.last_kwargs)
|
|
self.assertNotIn("initial_prompt", model.last_kwargs)
|
|
|
|
def test_transcribe_applies_hints_when_model_supports_them(self):
|
|
desktop = FakeDesktop()
|
|
model = FakeHintModel(text="hello")
|
|
cfg = self._config()
|
|
cfg.vocabulary.terms = ["Systemd"]
|
|
cfg.vocabulary.replacements = [VocabularyReplacement(source="docker", target="Docker")]
|
|
|
|
daemon = self._build_daemon(desktop, model, cfg=cfg, verbose=False)
|
|
|
|
result = daemon._transcribe(object())
|
|
|
|
self.assertEqual(result, "hello")
|
|
self.assertIn("Docker", model.last_kwargs["hotwords"])
|
|
self.assertIn("Systemd", model.last_kwargs["hotwords"])
|
|
self.assertIn("Preferred vocabulary", model.last_kwargs["initial_prompt"])
|
|
|
|
def test_verbose_flag_controls_transcript_logging(self):
|
|
desktop = FakeDesktop()
|
|
cfg = self._config()
|
|
|
|
daemon = self._build_daemon(desktop, FakeModel(), cfg=cfg, verbose=False)
|
|
self.assertFalse(daemon.log_transcript)
|
|
|
|
daemon_verbose = self._build_daemon(desktop, FakeModel(), cfg=cfg, verbose=True)
|
|
self.assertTrue(daemon_verbose.log_transcript)
|
|
|
|
def test_ai_processor_is_initialized_during_daemon_init(self):
|
|
desktop = FakeDesktop()
|
|
with patch("aman._build_whisper_model", return_value=FakeModel()), patch(
|
|
"aman.LlamaProcessor", return_value=FakeAIProcessor()
|
|
) as processor_cls:
|
|
daemon = aman.Daemon(self._config(), desktop, verbose=True)
|
|
|
|
processor_cls.assert_called_once_with(verbose=True)
|
|
self.assertIsNotNone(daemon.ai_processor)
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_passes_clipboard_remove_option_to_desktop(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
model = FakeModel(text="hello world")
|
|
cfg = self._config()
|
|
cfg.injection.remove_transcription_from_clipboard = True
|
|
|
|
daemon = self._build_daemon(desktop, model, cfg=cfg, verbose=False)
|
|
daemon._start_stop_worker = (
|
|
lambda stream, record, trigger, process_audio: daemon._stop_and_process(
|
|
stream, record, trigger, process_audio
|
|
)
|
|
)
|
|
|
|
daemon.toggle()
|
|
daemon.toggle()
|
|
|
|
self.assertEqual(desktop.inject_calls, [("hello world", "clipboard", True)])
|
|
|
|
def test_state_changes_are_debug_level(self):
|
|
desktop = FakeDesktop()
|
|
daemon = self._build_daemon(desktop, FakeModel(), verbose=False)
|
|
|
|
with self.assertLogs(level="DEBUG") as logs:
|
|
daemon.set_state(aman.State.RECORDING)
|
|
|
|
self.assertTrue(
|
|
any("DEBUG:root:state: idle -> recording" in line for line in logs.output)
|
|
)
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_escape_listener_is_only_armed_while_recording(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
daemon = self._build_daemon(desktop, FakeModel(), verbose=False)
|
|
daemon._start_stop_worker = (
|
|
lambda stream, record, trigger, process_audio: daemon._stop_and_process(
|
|
stream, record, trigger, process_audio
|
|
)
|
|
)
|
|
|
|
self.assertFalse(desktop.cancel_listener_active)
|
|
daemon.toggle()
|
|
self.assertTrue(desktop.cancel_listener_active)
|
|
daemon.toggle()
|
|
self.assertFalse(desktop.cancel_listener_active)
|
|
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_edit_mode_opens_popup_and_starts_recording(self, _start_mock):
|
|
desktop = FakeDesktop()
|
|
desktop.clipboard_text = "Hello team"
|
|
daemon = self._build_daemon(desktop, FakeModel(text="make it funnier"), verbose=False)
|
|
|
|
daemon.toggle_edit()
|
|
|
|
self.assertTrue(desktop.popup_open)
|
|
self.assertEqual(desktop.popup_text, "Hello team")
|
|
self.assertEqual(daemon.get_state(), aman.State.EDIT_RECORDING)
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_edit_mode_instruction_updates_popup_text(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
desktop.clipboard_text = "Hello team"
|
|
daemon = self._build_daemon(desktop, FakeModel(text="make it funnier"), verbose=False)
|
|
|
|
daemon.toggle_edit()
|
|
daemon.toggle_edit()
|
|
|
|
self.assertTrue(
|
|
self._wait_until(lambda: daemon.get_state() == aman.State.EDIT_IDLE),
|
|
"edit mode did not return to EDIT_IDLE",
|
|
)
|
|
self.assertEqual(desktop.popup_text, "Hello team [make it funnier]")
|
|
self.assertEqual(len(daemon.ai_processor.edit_calls), 1)
|
|
self.assertEqual(
|
|
daemon.ai_processor.edit_calls[0]["instruction_history"],
|
|
["make it funnier"],
|
|
)
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_enter_finalizes_and_injects(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
desktop.clipboard_text = "Initial"
|
|
daemon = self._build_daemon(desktop, FakeModel(text="instruction"), verbose=False)
|
|
|
|
daemon.toggle_edit()
|
|
desktop.popup_text = "Final text"
|
|
daemon.finalize_edit_session_inject()
|
|
|
|
self.assertTrue(
|
|
self._wait_until(lambda: len(desktop.inject_calls) == 1),
|
|
"edit finalize did not inject text",
|
|
)
|
|
self.assertFalse(desktop.popup_open)
|
|
self.assertEqual(desktop.inject_calls[0], ("Final text", "clipboard", False))
|
|
self.assertEqual(desktop.focus_restore_calls, 1)
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_ctrl_c_copies_and_closes_without_inject(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
desktop.clipboard_text = "Initial"
|
|
daemon = self._build_daemon(desktop, FakeModel(text="instruction"), verbose=False)
|
|
|
|
daemon.toggle_edit()
|
|
desktop.popup_text = "Copied text"
|
|
daemon.finalize_edit_session_copy()
|
|
|
|
self.assertTrue(
|
|
self._wait_until(lambda: not desktop.popup_open),
|
|
"edit popup did not close after copy",
|
|
)
|
|
self.assertEqual(desktop.clipboard_text, "Copied text")
|
|
self.assertEqual(desktop.inject_calls, [])
|
|
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_normal_hotkey_ignored_while_edit_session_active(self, _start_mock):
|
|
desktop = FakeDesktop()
|
|
desktop.clipboard_text = "Initial"
|
|
daemon = self._build_daemon(desktop, FakeModel(text="instruction"), verbose=False)
|
|
|
|
daemon.toggle_edit()
|
|
daemon.toggle()
|
|
|
|
self.assertEqual(daemon.get_state(), aman.State.EDIT_RECORDING)
|
|
|
|
@patch("aman.stop_audio_recording", return_value=FakeAudio(8))
|
|
@patch("aman.start_audio_recording", return_value=(object(), object()))
|
|
def test_handle_cancel_closes_edit_session(self, _start_mock, _stop_mock):
|
|
desktop = FakeDesktop()
|
|
desktop.clipboard_text = "Initial"
|
|
daemon = self._build_daemon(desktop, FakeModel(text="instruction"), verbose=False)
|
|
|
|
daemon.toggle_edit()
|
|
daemon.handle_cancel()
|
|
|
|
self.assertTrue(
|
|
self._wait_until(lambda: daemon.get_state() == aman.State.IDLE),
|
|
"edit cancel did not reach idle state",
|
|
)
|
|
self.assertFalse(desktop.popup_open)
|
|
|
|
|
|
class LockTests(unittest.TestCase):
|
|
def test_lock_rejects_second_instance(self):
|
|
with tempfile.TemporaryDirectory() as td:
|
|
with patch.dict(os.environ, {"XDG_RUNTIME_DIR": td}, clear=False):
|
|
first = aman._lock_single_instance()
|
|
try:
|
|
with self.assertRaises(SystemExit) as ctx:
|
|
aman._lock_single_instance()
|
|
self.assertIn("already running", str(ctx.exception))
|
|
finally:
|
|
first.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|