tapir/internal/uploads/worker_test.go

package uploads_test

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"git.d-ma.be/mathias/tapir/internal/domain"
	"git.d-ma.be/mathias/tapir/internal/uploads"
)

type wStore struct {
	mu        sync.Mutex
	started   []string
	completed map[string]domain.Transcription
	failed    map[string]domain.UploadErrorKind
	startErr  error
	// completeErr makes CompleteUpload fail, as a database outage would.
	completeErr error
	saves       [][]domain.ProgressStep
	beats       int
	// beatDelay makes each heartbeat write slow, so one is still in flight when a job
	// ends -- which is what proves the worker waits for its heartbeat to finish.
	beatDelay time.Duration
}

func newWStore() *wStore {
	return &wStore{completed: map[string]domain.Transcription{}, failed: map[string]domain.UploadErrorKind{}}
}

func (s *wStore) StartTranscribing(_ context.Context, _, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.started = append(s.started, id)
	return s.startErr
}

func (s *wStore) CompleteUpload(_ context.Context, _, id string, t domain.Transcription) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.completeErr != nil {
		return s.completeErr
	}
	s.completed[id] = t
	return nil
}

func (s *wStore) FailUpload(_ context.Context, _, id string, k domain.UploadErrorKind) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.failed[id] = k
	return nil
}

func (s *wStore) SaveProgress(_ context.Context, _, _ string, steps []domain.ProgressStep, _ time.Time) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.saves = append(s.saves, append([]domain.ProgressStep(nil), steps...))
	return nil
}

func (s *wStore) Heartbeat(context.Context, string, string) error {
	time.Sleep(s.beatDelay)
	s.mu.Lock()
	defer s.mu.Unlock()
	s.beats++
	return nil
}

func (s *wStore) lastSteps() []domain.ProgressStep {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.saves) == 0 {
		return nil
	}
	return s.saves[len(s.saves)-1]
}

func (s *wStore) beatCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.beats
}

func (s *wStore) snapshot() (map[string]domain.Transcription, map[string]domain.UploadErrorKind) {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := map[string]domain.Transcription{}
	f := map[string]domain.UploadErrorKind{}
	for k, v := range s.completed {
		c[k] = v
	}
	for k, v := range s.failed {
		f[k] = v
	}
	return c, f
}

type fakeConverter struct {
	err  error
	srcs []string
	dsts []string
}

func (c *fakeConverter) ToWAV(_ context.Context, src, dst string) error {
	c.srcs = append(c.srcs, src)
	c.dsts = append(c.dsts, dst)
	if c.err != nil {
		return c.err
	}
	return os.WriteFile(dst, []byte("wav"), 0o600)
}

type fakeTranscriber struct {
	result domain.Transcription
	err    error
	paths  []string
	// block, when non-nil, holds every call until it is closed or ctx ends.
	block   chan struct{}
	running atomic.Int32
	maxSeen atomic.Int32
}

func (f *fakeTranscriber) Transcribe(ctx context.Context, wav string) (domain.Transcription, error) {
	f.paths = append(f.paths, wav)
	n := f.running.Add(1)
	defer f.running.Add(-1)
	for {
		m := f.maxSeen.Load()
		if n <= m || f.maxSeen.CompareAndSwap(m, n) {
			break
		}
	}
	if f.block != nil {
		select {
		case <-f.block:
		case <-ctx.Done():
			return domain.Transcription{}, ctx.Err()
		}
	}
	return f.result, f.err
}

type wrig struct {
	w    *uploads.Worker
	st   *wStore
	conv *fakeConverter
	tr   *fakeTranscriber
	tmp  string
}

func newWRig(t *testing.T, timeout time.Duration) *wrig {
	t.Helper()
	r := &wrig{
		st:   newWStore(),
		conv: &fakeConverter{},
		tr:   &fakeTranscriber{result: domain.Transcription{Text: "synthetic text", Endpoint: "iguana/kb-whisper", AudioSeconds: 3}},
		tmp:  t.TempDir(),
	}
	r.w = uploads.NewWorker(r.st, r.conv, r.tr, r.tmp, timeout)
	return r
}

// spool writes a fake uploaded file (what Intake would hand over) and returns a Job.
func (r *wrig) spool(t *testing.T, id string) uploads.Job {
	t.Helper()
	p := filepath.Join(r.tmp, "upload-"+id)
	require.NoError(t, os.WriteFile(p, []byte("audio"), 0o600))
	return uploads.Job{UserID: "u1", UploadID: id, Path: p}
}

func (r *wrig) leftovers(t *testing.T) []string {
	t.Helper()
	ents, err := os.ReadDir(r.tmp)
	require.NoError(t, err)
	var n []string
	for _, e := range ents {
		n = append(n, e.Name())
	}
	return n
}

func TestWorkerProcess_SuccessConvertsTranscribesStoresAndDeletesTheAudio(t *testing.T) {
	r := newWRig(t, time.Minute)
	job := r.spool(t, "a")

	r.w.Process(context.Background(), job)

	c, f := r.st.snapshot()
	assert.Equal(t, []string{"a"}, r.st.started)
	assert.Equal(t, "synthetic text", c["a"].Text)
	assert.Equal(t, "iguana/kb-whisper", c["a"].Endpoint)
	assert.Empty(t, f)
	require.Len(t, r.conv.srcs, 1)
	assert.Equal(t, job.Path, r.conv.srcs[0], "the converter reads the uploaded file")
	assert.Equal(t, r.conv.dsts, r.tr.paths, "whisper is sent the converted WAV, not the original")
	assert.Equal(t, r.tmp, filepath.Dir(r.conv.dsts[0]))
	assert.Empty(t, r.leftovers(t), "the audio is never retained: both the upload and the WAV are deleted")
}

func TestWorkerProcess_FailuresAreTypedAndStillDeleteTheAudio(t *testing.T) {
	tests := []struct {
		name  string
		setup func(*wrig)
	}{
		{"conversion fails", func(r *wrig) { r.conv.err = errors.New("moov atom not found") }},
		{"every transcription endpoint fails", func(r *wrig) { r.tr.err = errors.New("iguana down; berget down") }},
		{"storing the finished transcript fails", func(r *wrig) { r.st.completeErr = errors.New("db down") }},
		{"the temporary WAV cannot be created", func(r *wrig) {
			r.w = uploads.NewWorker(r.st, r.conv, r.tr, filepath.Join(r.tmp, "missing"), time.Minute)
		}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			r := newWRig(t, time.Minute)
			tt.setup(r)
			job := r.spool(t, "a")

			r.w.Process(context.Background(), job)

			c, f := r.st.snapshot()
			assert.Empty(t, c)
			assert.Equal(t, domain.UploadErrTranscriptionFailed, f["a"])
			assert.Empty(t, r.leftovers(t))
		})
	}
	t.Run("conversion failure never reaches whisper", func(t *testing.T) {
		r := newWRig(t, time.Minute)
		r.conv.err = errors.New("bad")

		r.w.Process(context.Background(), r.spool(t, "a"))

		assert.Empty(t, r.tr.paths)
	})
}

func TestWorkerProcess_ShutdownMidJobLeavesTheRowForTheStartupSweep(t *testing.T) {
	r := newWRig(t, time.Minute)
	r.tr.block = make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	job := r.spool(t, "a")

	go func() { r.w.Process(ctx, job); close(done) }()
	require.Eventually(t, func() bool { return r.tr.running.Load() == 1 }, 2*time.Second, 5*time.Millisecond)
	cancel()
	<-done

	_, f := r.st.snapshot()
	assert.Empty(t, f, "a shutdown is not a transcription failure: the row stays for FailInterruptedUploads to mark interrupted")
	assert.Empty(t, r.leftovers(t), "but the audio is still deleted")
}

func TestWorkerProcess_AJobThatOutlivesItsTimeoutIsFailedNotStuck(t *testing.T) {
	r := newWRig(t, 50*time.Millisecond)
	r.tr.block = make(chan struct{}) // never released
	done := make(chan struct{})

	go func() { r.w.Process(context.Background(), r.spool(t, "a")); close(done) }()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("Process did not return: a hung endpoint was not bounded by the job timeout")
	}

	_, f := r.st.snapshot()
	assert.Equal(t, domain.UploadErrTranscriptionFailed, f["a"], "a hung endpoint must end in failed, not an eternal spinner (ADR-025)")
	assert.Empty(t, r.leftovers(t))
}

func TestWorkerProcess_StartFailureDropsTheJobAndDeletesTheAudio(t *testing.T) {
	r := newWRig(t, time.Minute)
	r.st.startErr = errors.New("db down")

	r.w.Process(context.Background(), r.spool(t, "a"))

	assert.Empty(t, r.tr.paths, "never transcribe a job whose row could not be moved to transcribing")
	_, f := r.st.snapshot()
	assert.Contains(t, f, "a")
	assert.Empty(t, r.leftovers(t))
}

func TestWorkerRun_ProcessesOneJobAtATime(t *testing.T) {
	r := newWRig(t, time.Minute)
	r.tr.block = make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go r.w.Run(ctx)

	for _, id := range []string{"a", "b", "c"} {
		require.NoError(t, r.w.Submit(r.spool(t, id)))
	}
	require.Eventually(t, func() bool { return r.tr.running.Load() == 1 }, 2*time.Second, 5*time.Millisecond)
	time.Sleep(50 * time.Millisecond) // give a (wrong) second job time to start
	assert.Equal(t, int32(1), r.tr.running.Load(), "two simultaneous long jobs on iguana were never measured: concurrency is 1")

	close(r.tr.block)
	require.Eventually(t, func() bool { c, _ := r.st.snapshot(); return len(c) == 3 }, 3*time.Second, 5*time.Millisecond)
	assert.Equal(t, int32(1), r.tr.maxSeen.Load())
}

func TestWorkerRun_StoppingDeletesQueuedAudioAndRefusesNewJobs(t *testing.T) {
	r := newWRig(t, time.Minute)
	r.tr.block = make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	stopped := make(chan struct{})
	go func() { r.w.Run(ctx); close(stopped) }()
	require.NoError(t, r.w.Submit(r.spool(t, "a")))
	require.Eventually(t, func() bool { return r.tr.running.Load() == 1 }, 2*time.Second, 5*time.Millisecond)
	require.NoError(t, r.w.Submit(r.spool(t, "b"))) // queued behind the running job

	cancel()
	<-stopped

	assert.Len(t, r.conv.srcs, 1, "a job still queued at shutdown must never be started: only the running job was converted")
	assert.Empty(t, r.leftovers(t), "audio of queued jobs is deleted on stop: the rows are swept to interrupted at next start")
	assert.ErrorIs(t, r.w.Submit(r.spool(t, "c")), uploads.ErrStopped)
}