stt/stress_recording.py at master · bokan/stt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#!/usr/bin/env python3
"""
Bruteforce audio recording start/stop/cancel to surface stuck states.
"""

from __future__ import annotations

import argparse
import os
import random
import sys
import threading
import time

os.environ.setdefault("STT_HEADLESS", "1")

import stt


class _NullOverlay:
    def show(self):
        pass

    def hide(self):
        pass

    def update_waveform(self, values, above_threshold=False):
        pass

    def set_transcribing(self, transcribing: bool):
        pass

    def set_shift_held(self, held: bool):
        pass


class _DummyProvider:
    name = "dummy"

    def is_available(self):
        return True

    def warmup(self):
        return None

    def transcribe(self, audio_file_path: str, language: str, prompt: str | None = None):
        return ""


def _safe_unlink(path: str | None) -> None:
    if not path:
        return
    try:
        os.unlink(path)
    except OSError:
        pass


def main() -> int:
    parser = argparse.ArgumentParser(description="Stress test STT recording engine")
    parser.add_argument("--iterations", type=int, default=200)
    parser.add_argument("--min-hold", type=float, default=0.02)
    parser.add_argument("--max-hold", type=float, default=0.5)
    parser.add_argument("--min-gap", type=float, default=0.01)
    parser.add_argument("--max-gap", type=float, default=0.2)
    parser.add_argument("--mode", choices=["stop", "cancel", "mixed", "full"], default="mixed")
    parser.add_argument("--cancel-rate", type=float, default=0.2)
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--async-start", action="store_true")
    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)

    app = stt.STTApp(
        device_name=args.device,
        provider=_DummyProvider(),
        overlay=_NullOverlay(),
        sound_player=lambda _: None,
        text_injector=lambda *_: None,
        keep_recordings=False,
    )

    if args.mode == "full":
        app.type_text = lambda *_, **__: None
        app.print_ready_prompt = lambda: None

    stats = {
        "iterations": 0,
        "stops": 0,
        "cancels": 0,
        "errors": 0,
        "empty_wav": 0,
    }

    try:
        for i in range(args.iterations):
            stats["iterations"] += 1

            def _start():
                app.start_recording()

            if args.async_start:
                threading.Thread(target=_start, daemon=True).start()
            else:
                _start()

            time.sleep(random.uniform(args.min_hold, args.max_hold))

            action = args.mode
            if args.mode == "mixed":
                action = "cancel" if random.random() < args.cancel_rate else "stop"

            try:
                if action == "cancel":
                    app.cancel_recording()
                    stats["cancels"] += 1
                elif action == "full":
                    app.process_recording(send_enter=False)
                    stats["stops"] += 1
                else:
                    wav_path, frames, peak = app.stop_recording()
                    if not wav_path:
                        stats["empty_wav"] += 1
                    _safe_unlink(wav_path)
                    stats["stops"] += 1
            except Exception:
                stats["errors"] += 1

            time.sleep(random.uniform(args.min_gap, args.max_gap))
    finally:
        try:
            app._audio_worker.stop(force=True)
        except Exception:
            pass

    print("stress_done", stats)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())