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
}

func (f *fakeTranscriber) Decode([]float32) string {
	if len(f.decodeTexts) == 0 {
		return f.text
	}
	text := f.decodeTexts[0]
	f.decodeTexts = f.decodeTexts[1:]
	return text
}
func (f *fakeTranscriber) AcceptVAD([]float32) {}
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}, 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(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}, transcriber, func(event Event) error {
		events = append(events, event)
		return nil
	})

	if err := processor.Accept(make([]float32, audio.SampleRate+audio.SampleRate/2)); 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"}
	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 := make([]float32, audio.SampleRate)
	for i := range loud {
		loud[i] = 0.5
	}
	if err := processor.Accept(loud); err != nil {
		t.Fatal(err)
	}
	transcriber.segments = []fakeSegment{{start: 0, samples: loud}}
	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}, transcriber, func(got Event) error {
		event = got
		return nil
	})
	if err := processor.Accept(make([]float32, audio.SampleRate)); err != nil {
		t.Fatal(err)
	}
	if event.Kind != Hide {
		t.Fatalf("event kind = %d, want Hide", event.Kind)
	}
}

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")
	}
}