gpui-native-kit/crates/core/tests/speech.rs

#![cfg(feature = "speech")]

use std::cell::RefCell;
use std::rc::Rc;
use std::time::Duration;

use gpui::{App, AppContext as _, Entity, SharedString, Subscription, TestAppContext};
use gpui_native_core::speech::{
    AudioFormat, AudioInput, AudioSink, RecognitionSession, SpeechError, SpeechEvent,
    SpeechRecognizer, SpeechSink, SpeechState, SpeechStatus,
};

type EventLog = Rc<RefCell<Vec<SpeechEvent>>>;
type Harness = (
    Entity<SpeechState>,
    EventLog,
    Rc<RefCell<Option<SpeechSink>>>,
    Rc<RefCell<usize>>,
);

struct FakeRecognizer {
    available: bool,
    finish_reports: bool,
    sink: Rc<RefCell<Option<SpeechSink>>>,
    received: Rc<RefCell<usize>>,
}

impl SpeechRecognizer for FakeRecognizer {
    fn is_available(&self, _cx: &App) -> bool {
        self.available
    }

    fn start(
        &self,
        sink: SpeechSink,
        cx: &mut App,
    ) -> Result<Box<dyn RecognitionSession>, SpeechError> {
        *self.sink.borrow_mut() = Some(sink.clone());
        sink.ready(cx);
        Ok(Box::new(FakeSession {
            sink,
            finish_reports: self.finish_reports,
            received: self.received.clone(),
        }))
    }
}

struct FakeSession {
    sink: SpeechSink,
    finish_reports: bool,
    received: Rc<RefCell<usize>>,
}

impl RecognitionSession for FakeSession {
    fn push_audio(&mut self, samples: &[i16], _cx: &mut App) {
        *self.received.borrow_mut() += samples.len();
    }

    fn finish(&mut self, cx: &mut App) {
        if self.finish_reports {
            self.sink.finish(cx);
        }
    }
}

/// Delivers one silent second, 16 kHz mono, in 100 ms chunks.
struct FakeInput;

impl AudioInput for FakeInput {
    fn start(
        &self,
        format: AudioFormat,
        sink: AudioSink,
        cx: &mut App,
    ) -> Result<Subscription, SpeechError> {
        let chunk_len = (format.sample_rate() / 10) as usize * format.channels() as usize;
        let task = cx.spawn(async move |cx| {
            for _ in 0..10 {
                cx.background_executor()
                    .timer(Duration::from_millis(100))
                    .await;
                cx.update(|cx| sink.push(vec![0; chunk_len], cx));
            }
        });
        Ok(Subscription::new(move || drop(task)))
    }
}

fn harness(cx: &mut TestAppContext, available: bool, finish_reports: bool) -> Harness {
    let sink_slot: Rc<RefCell<Option<SpeechSink>>> = Rc::new(RefCell::new(None));
    let received = Rc::new(RefCell::new(0));
    let log: EventLog = Rc::new(RefCell::new(Vec::new()));
    let record = log.clone();
    let speech = cx.update(|cx| {
        let speech = cx.new(|cx| {
            SpeechState::new(cx)
                .recognizer(FakeRecognizer {
                    available,
                    finish_reports,
                    sink: sink_slot.clone(),
                    received: received.clone(),
                })
                .input(FakeInput)
                .stop_timeout(Duration::from_secs(2))
        });
        cx.subscribe(&speech, move |_, event, _| {
            record.borrow_mut().push(event.clone())
        })
        .detach();
        speech
    });
    (speech, log, sink_slot, received)
}

fn start(speech: &Entity<SpeechState>, cx: &mut TestAppContext) {
    cx.update(|cx| speech.update(cx, |state, cx| state.start(cx)));
    cx.run_until_parked();
}

#[test]
fn a_session_starts_reads_and_finalizes() {
    let mut cx = TestAppContext::single();
    let (speech, events, sink_slot, _) = harness(&mut cx, true, true);

    start(&speech, &mut cx);
    cx.update(|cx| assert_eq!(speech.read(cx).status(), SpeechStatus::Recording));
    assert_eq!(*events.borrow(), vec![SpeechEvent::Started]);

    let sink = sink_slot.borrow().clone().expect("session opened");
    cx.update(|cx| sink.hypothesis("hel", cx));
    cx.run_until_parked();
    cx.update(|cx| {
        sink.phrase("hello ", cx);
        sink.hypothesis("wor", cx);
    });
    cx.run_until_parked();
    cx.update(|cx| {
        assert_eq!(
            speech.read(cx).transcript(),
            SharedString::from("hello wor")
        );
    });

    cx.update(|cx| speech.update(cx, |state, cx| state.stop(cx)));
    cx.run_until_parked();
    cx.update(|cx| assert_eq!(speech.read(cx).status(), SpeechStatus::Idle));

    let events = events.borrow();
    assert!(
        matches!(events.last(), Some(SpeechEvent::Final(text)) if text.as_ref() == "hello wor")
    );
    assert!(matches!(
        events.get(events.len() - 2),
        Some(SpeechEvent::Partial(text)) if text.as_ref() == "hello wor"
    ));
}

#[test]
fn starting_twice_runs_a_single_session() {
    let mut cx = TestAppContext::single();
    let (speech, events, _, _) = harness(&mut cx, true, true);

    start(&speech, &mut cx);
    start(&speech, &mut cx);

    assert_eq!(*events.borrow(), vec![SpeechEvent::Started]);
}

#[test]
fn cancel_discards_the_transcript_and_ignores_late_reports() {
    let mut cx = TestAppContext::single();
    let (speech, events, sink_slot, _) = harness(&mut cx, true, true);

    start(&speech, &mut cx);
    let sink = sink_slot.borrow().clone().expect("session opened");
    cx.update(|cx| sink.phrase("hello", cx));
    cx.run_until_parked();

    cx.update(|cx| speech.update(cx, |state, cx| state.cancel(cx)));
    cx.run_until_parked();
    cx.update(|cx| {
        let state = speech.read(cx);
        assert_eq!(state.status(), SpeechStatus::Idle);
        assert_eq!(state.transcript(), SharedString::from(""));
    });
    assert!(matches!(
        events.borrow().last(),
        Some(SpeechEvent::Cancelled)
    ));

    cx.update(|cx| sink.finish(cx));
    cx.run_until_parked();
    assert!(
        !events
            .borrow()
            .iter()
            .any(|event| matches!(event, SpeechEvent::Final(_)))
    );
}

#[test]
fn starting_without_a_recognizer_fails() {
    let mut cx = TestAppContext::single();
    let log: EventLog = Rc::new(RefCell::new(Vec::new()));
    let record = log.clone();
    let speech = cx.update(|cx| {
        let speech = cx.new(|cx| SpeechState::new(cx).input(FakeInput));
        cx.subscribe(&speech, move |_, event, _| {
            record.borrow_mut().push(event.clone())
        })
        .detach();
        speech
    });

    cx.update(|cx| assert!(!speech.read(cx).has_recognizer()));
    start(&speech, &mut cx);

    cx.update(|cx| assert_eq!(speech.read(cx).status(), SpeechStatus::Idle));
    assert!(matches!(events_last(&log), Some(SpeechEvent::Error(_))));
}

#[test]
fn starting_an_unavailable_recognizer_fails() {
    let mut cx = TestAppContext::single();
    let (speech, events, _, _) = harness(&mut cx, false, true);

    start(&speech, &mut cx);
    cx.update(|cx| assert_eq!(speech.read(cx).status(), SpeechStatus::Idle));
    assert!(matches!(
        events.borrow().last(),
        Some(SpeechEvent::Error(_))
    ));
}

#[test]
fn the_stop_timeout_finalizes_with_the_transcript_so_far() {
    let mut cx = TestAppContext::single();
    let (speech, events, sink_slot, _) = harness(&mut cx, true, false);

    start(&speech, &mut cx);
    let sink = sink_slot.borrow().clone().expect("session opened");
    cx.update(|cx| sink.phrase("hello", cx));
    cx.run_until_parked();

    cx.update(|cx| speech.update(cx, |state, cx| state.stop(cx)));
    cx.run_until_parked();
    cx.update(|cx| assert_eq!(speech.read(cx).status(), SpeechStatus::Stopping));

    cx.executor().advance_clock(Duration::from_secs(2));
    cx.run_until_parked();
    cx.update(|cx| assert_eq!(speech.read(cx).status(), SpeechStatus::Idle));
    assert!(
        matches!(events.borrow().last(), Some(SpeechEvent::Final(text)) if text.as_ref() == "hello")
    );
}

#[test]
fn captured_audio_reaches_the_session_and_the_levels() {
    let mut cx = TestAppContext::single();
    let (speech, _, _, received) = harness(&mut cx, true, true);

    start(&speech, &mut cx);
    cx.executor().advance_clock(Duration::from_secs(1));
    cx.run_until_parked();

    cx.update(|cx| {
        let state = speech.read(cx);
        assert_eq!(*received.borrow(), 16_000);
        // One level per 80 ms bucket of the one-second capture.
        let levels = state.levels();
        assert_eq!(levels.len(), 12);
        assert!(levels.iter().all(|level| *level == 0.0));
    });

    cx.update(|cx| speech.update(cx, |state, cx| state.stop(cx)));
    cx.run_until_parked();
    cx.update(|cx| assert!(speech.read(cx).levels().is_empty()));
}

fn events_last(log: &EventLog) -> Option<SpeechEvent> {
    log.borrow().last().cloned()
}