Captioneer/src/captions/processor_test.go

package captions

import (
	"testing"
	"time"

	"tea.chunkbyte.com/kato/captioneer/src/audio"
	"tea.chunkbyte.com/kato/captioneer/src/config"
)

type fakeSegment struct {
	start   int
	samples []float32
}

type fakeTranscriber struct {
	text         string
	decodeTexts  []string
	segments     []fakeSegment
	flushed      bool
	speechActive bool
	decodeSizes  []int
	acceptedVAD  int
}

func (f *fakeTranscriber) Decode(samples []float32) string {
	f.decodeSizes = append(f.decodeSizes, len(samples))
	if len(f.decodeTexts) == 0 {
		return f.text
	}
	text := f.decodeTexts[0]
	f.decodeTexts = f.decodeTexts[1:]
	return text
}
func (f *fakeTranscriber) AcceptVAD(samples []float32) { f.acceptedVAD += len(samples) }
func (f *fakeTranscriber) SpeechActive() bool          { return f.speechActive }
func (f *fakeTranscriber) FlushVAD()                   { f.flushed = true }
func (f *fakeTranscriber) NextSpeechSegment() (int, []float32, bool) {
	if len(f.segments) == 0 {
		return 0, nil, false
	}
	segment := f.segments[0]
	f.segments = f.segments[1:]
	return segment.start, segment.samples, true
}

func TestFixedModeEmitsFinalCaption(t *testing.T) {
	transcriber := &fakeTranscriber{text: "hello"}
	var events []Event
	processor := New(config.Settings{
		Mode: config.ModeFixed, ChunkDuration: time.Second, PreviewThresholdDB: -45,
	}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	if err := processor.Accept(constantSamples(audio.SampleRate, 0.5)); err != nil {
		t.Fatal(err)
	}
	if len(events) != 1 || events[0].Kind != Final || events[0].Text != "hello" {
		t.Fatalf("unexpected events: %#v", events)
	}
	if events[0].StartedAt != 0 || events[0].EndedAt != time.Second {
		t.Fatalf("unexpected timestamps: %#v", events[0])
	}
}

func TestFixedModeFlushesFinalPartialChunkWithExactTimestamps(t *testing.T) {
	transcriber := &fakeTranscriber{decodeTexts: []string{"first", "partial"}}
	var events []Event
	processor := New(config.Settings{
		Mode: config.ModeFixed, ChunkDuration: time.Second, PreviewThresholdDB: -45,
	}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	if err := processor.Accept(constantSamples(audio.SampleRate+audio.SampleRate/2, 0.5)); err != nil {
		t.Fatal(err)
	}
	if err := processor.Flush(); err != nil {
		t.Fatal(err)
	}
	if len(events) != 2 {
		t.Fatalf("event count = %d, want 2", len(events))
	}
	if events[1].Text != "partial" || events[1].StartedAt != time.Second || events[1].EndedAt != 1500*time.Millisecond {
		t.Fatalf("unexpected partial event: %#v", events[1])
	}
}

func TestVADModeEmitsProvisionalThenFinal(t *testing.T) {
	transcriber := &fakeTranscriber{text: "speech", speechActive: true}
	var events []Event
	processor := New(config.Settings{
		Mode:               config.ModeVAD,
		ChunkDuration:      time.Second,
		PreviewThresholdDB: -45,
	}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	loud := constantSamples(audio.SampleRate, 0.5)
	if err := processor.Accept(loud); err != nil {
		t.Fatal(err)
	}
	transcriber.segments = []fakeSegment{{start: 0, samples: loud}}
	transcriber.speechActive = false
	if err := processor.Accept(make([]float32, audio.SamplesPerPacket())); err != nil {
		t.Fatal(err)
	}

	if len(events) != 2 || events[0].Kind != Provisional || events[1].Kind != Final {
		t.Fatalf("unexpected events: %#v", events)
	}
}

func TestEmptyFinalEmitsHide(t *testing.T) {
	transcriber := &fakeTranscriber{text: " \n\t "}
	var event Event
	processor := New(config.Settings{
		Mode: config.ModeFixed, ChunkDuration: time.Second, PreviewThresholdDB: -45,
	}, transcriber, func(got Event) error {
		event = got
		return nil
	})
	if err := processor.Accept(constantSamples(audio.SampleRate, 0.5)); err != nil {
		t.Fatal(err)
	}
	if event.Kind != Hide {
		t.Fatalf("event kind = %d, want Hide", event.Kind)
	}
}

func TestFixedModeSilenceStaysOutOfRecognizer(t *testing.T) {
	transcriber := &fakeTranscriber{text: "hallucinated caption"}
	var events []Event
	processor := New(config.Settings{
		Mode:               config.ModeFixed,
		ChunkDuration:      time.Second,
		PreviewThresholdDB: -45,
	}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	if err := processor.Accept(make([]float32, audio.SampleRate)); err != nil {
		t.Fatal(err)
	}
	if len(transcriber.decodeSizes) != 0 {
		t.Fatalf("silent audio reached recognizer: decode sizes = %v", transcriber.decodeSizes)
	}
	if len(events) != 1 || events[0].Kind != Hide {
		t.Fatalf("silent chunk events = %#v, want one Hide", events)
	}
}

func TestVADModeDoesNotPreviewLoudNonSpeech(t *testing.T) {
	transcriber := &fakeTranscriber{text: "music hallucination"}
	var events []Event
	processor := New(config.Settings{
		Mode:               config.ModeVAD,
		ChunkDuration:      time.Second,
		PreviewThresholdDB: -45,
	}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	loudMusic := constantSamples(audio.SampleRate, 0.5)
	if err := processor.Accept(loudMusic); err != nil {
		t.Fatal(err)
	}
	if transcriber.acceptedVAD != len(loudMusic) {
		t.Fatalf("VAD received %d samples, want %d", transcriber.acceptedVAD, len(loudMusic))
	}
	if len(transcriber.decodeSizes) != 0 || len(events) != 0 {
		t.Fatalf("non-speech reached recognizer: decodes=%v events=%#v", transcriber.decodeSizes, events)
	}
}

func TestVADModeClearsRejectedShortSpeechPreview(t *testing.T) {
	transcriber := &fakeTranscriber{text: "draft", speechActive: true}
	var events []Event
	processor := New(config.Settings{
		Mode:               config.ModeVAD,
		ChunkDuration:      time.Second,
		PreviewThresholdDB: -45,
	}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	if err := processor.Accept(constantSamples(audio.SampleRate, 0.5)); err != nil {
		t.Fatal(err)
	}
	transcriber.speechActive = false
	if err := processor.Accept(make([]float32, audio.SamplesPerPacket())); err != nil {
		t.Fatal(err)
	}
	if len(events) != 2 || events[0].Kind != Provisional || events[1].Kind != Hide {
		t.Fatalf("events = %#v, want Provisional then Hide", events)
	}
}

func TestVADPreviewRecognitionWindowIsBounded(t *testing.T) {
	transcriber := &fakeTranscriber{text: "draft", speechActive: true}
	processor := New(config.Settings{
		Mode:               config.ModeVAD,
		ChunkDuration:      time.Second,
		PreviewThresholdDB: -45,
	}, transcriber, func(Event) error { return nil })

	for range 15 {
		if err := processor.Accept(constantSamples(audio.SampleRate, 0.5)); err != nil {
			t.Fatal(err)
		}
	}
	maxSamples := audio.DurationSamples(maxPreviewDuration)
	for _, size := range transcriber.decodeSizes {
		if size > maxSamples {
			t.Fatalf("preview decode used %d samples, limit is %d", size, maxSamples)
		}
	}
	if len(transcriber.decodeSizes) != 15 {
		t.Fatalf("preview decode count = %d, want 15", len(transcriber.decodeSizes))
	}
}

func TestFlushFlushesVAD(t *testing.T) {
	transcriber := &fakeTranscriber{}
	processor := New(config.Settings{Mode: config.ModeVAD, ChunkDuration: time.Second}, transcriber, func(Event) error { return nil })
	if err := processor.Flush(); err != nil {
		t.Fatal(err)
	}
	if !transcriber.flushed {
		t.Fatal("VAD was not flushed")
	}
}

func constantSamples(count int, value float32) []float32 {
	samples := make([]float32, count)
	for i := range samples {
		samples[i] = value
	}
	return samples
}