coder/coderd/x/chatd/generation_retry_internal_test.go

package chatd //nolint:testpackage // Exercises unexported generation retry helpers.

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/testutil"
	"github.com/coder/quartz"
)

func TestTerminalGeneration(t *testing.T) {
	t.Parallel()

	require.Nil(t, terminalGeneration(nil))

	cause := xerrors.New("boom")
	wrapped := terminalGeneration(cause)
	require.True(t, isTerminalGeneration(wrapped))
	require.ErrorIs(t, wrapped, cause)
	require.ErrorIs(t, wrapped, errTerminalGeneration)
	require.Equal(t, cause.Error(), wrapped.Error())

	require.False(t, isTerminalGeneration(cause))
	require.False(t, isTerminalGeneration(nil))
}

func TestGenerationPhaseBackoff(t *testing.T) {
	t.Parallel()

	require.Equal(t, generationPhaseBaseBackoff, generationPhaseBackoff(0))
	require.Equal(t, 2*generationPhaseBaseBackoff, generationPhaseBackoff(1))
	require.Equal(t, 4*generationPhaseBaseBackoff, generationPhaseBackoff(2))
}

func TestRetryGenerationPhase(t *testing.T) {
	t.Parallel()

	t.Run("SuccessFirstTry", func(t *testing.T) {
		t.Parallel()
		starter := newGenerationPhaseTestStarter(t, quartz.NewMock(t).WithLogger(quartz.NoOpLogger))
		calls := 0
		got, err := retryGenerationPhase(context.Background(), starter, "prepare", func() (int, error) {
			calls++
			return 42, nil
		})
		require.NoError(t, err)
		require.Equal(t, 42, got)
		require.Equal(t, 1, calls)
	})

	t.Run("RetryThenSuccess", func(t *testing.T) {
		t.Parallel()
		clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger)
		timerTrap := clock.Trap().NewTimer("chatworker", "generation-phase-retry")
		defer timerTrap.Close()
		sink := testutil.NewFakeSink(t)
		starter := newGenerationPhaseTestStarter(t, clock)
		starter.opts.Logger = sink.Logger()
		ctx := testutil.Context(t, testutil.WaitLong)
		calls := 0
		done := make(chan phaseRetryResult[string], 1)
		go func() {
			got, err := retryGenerationPhase(ctx, starter, "prepare", func() (string, error) {
				calls++
				if calls < 2 {
					return "", xerrors.New("transient")
				}
				return "ok", nil
			})
			done <- phaseRetryResult[string]{value: got, err: err}
		}()

		timerTrap.MustWait(ctx).MustRelease(ctx)
		clock.Advance(generationPhaseBackoff(0)).MustWait(ctx)
		result := <-done
		require.NoError(t, result.err)
		require.Equal(t, "ok", result.value)
		require.Equal(t, 2, calls)
		entries := entriesWithMessage(sink, "chat generation phase retrying")
		require.Len(t, entries, 1)
		require.Equal(t, "prepare", sinkFieldValue(t, entries[0].Fields, "phase"))
		require.Equal(t, "1", sinkFieldValue(t, entries[0].Fields, "attempt"))
		require.Equal(t, generationPhaseBackoff(0).String(), sinkFieldValue(t, entries[0].Fields, "delay"))
	})

	t.Run("ExhaustsAndReturnsLastError", func(t *testing.T) {
		t.Parallel()
		clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger)
		timerTrap := clock.Trap().NewTimer("chatworker", "generation-phase-retry")
		defer timerTrap.Close()
		starter := newGenerationPhaseTestStarter(t, clock)
		ctx := testutil.Context(t, testutil.WaitLong)
		calls := 0
		done := make(chan error, 1)
		go func() {
			_, err := retryGenerationPhase(ctx, starter, "prepare", func() (int, error) {
				calls++
				return 0, xerrors.Errorf("attempt %d", calls)
			})
			done <- err
		}()

		timerTrap.MustWait(ctx).MustRelease(ctx)
		clock.Advance(generationPhaseBackoff(0)).MustWait(ctx)
		timerTrap.MustWait(ctx).MustRelease(ctx)
		clock.Advance(generationPhaseBackoff(1)).MustWait(ctx)
		require.EqualError(t, <-done, "attempt 3")
		require.Equal(t, generationPhaseMaxAttempts, calls)
	})

	t.Run("TerminalShortCircuits", func(t *testing.T) {
		t.Parallel()
		starter := newGenerationPhaseTestStarter(t, quartz.NewMock(t).WithLogger(quartz.NoOpLogger))
		calls := 0
		cause := xerrors.New("deterministic")
		_, err := retryGenerationPhase(context.Background(), starter, "prepare", func() (int, error) {
			calls++
			return 0, terminalGeneration(cause)
		})
		require.ErrorIs(t, err, cause)
		require.True(t, isTerminalGeneration(err))
		require.Equal(t, 1, calls)
	})

	t.Run("ContextCanceledExitsCleanly", func(t *testing.T) {
		t.Parallel()
		starter := newGenerationPhaseTestStarter(t, quartz.NewMock(t).WithLogger(quartz.NoOpLogger))
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		calls := 0
		_, err := retryGenerationPhase(ctx, starter, "prepare", func() (int, error) {
			calls++
			return 0, xerrors.New("transient")
		})
		require.ErrorIs(t, err, errTaskExpectedExit)
		require.Equal(t, 1, calls)
	})

	t.Run("WaitCancellationExitsCleanly", func(t *testing.T) {
		t.Parallel()
		clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger)
		timerTrap := clock.Trap().NewTimer("chatworker", "generation-phase-retry")
		defer timerTrap.Close()
		starter := newGenerationPhaseTestStarter(t, clock)
		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitLong))
		calls := 0
		done := make(chan error, 1)
		go func() {
			_, err := retryGenerationPhase(ctx, starter, "prepare", func() (int, error) {
				calls++
				return 0, xerrors.New("transient")
			})
			done <- err
		}()

		timerTrap.MustWait(ctx).MustRelease(ctx)
		cancel()
		require.ErrorIs(t, <-done, errTaskExpectedExit)
		require.Equal(t, 1, calls)
	})
}

type phaseRetryResult[T any] struct {
	value T
	err   error
}

func newGenerationPhaseTestStarter(t *testing.T, clock quartz.Clock) *taskStarter {
	t.Helper()
	require.NotNil(t, clock)
	return &taskStarter{opts: chatWorkerOptions{
		Clock:                   clock,
		Logger:                  testutil.NewFakeSink(t).Logger(),
		TaskRetryInitialBackoff: time.Millisecond,
		TaskRetryMaxBackoff:     time.Millisecond,
	}}
}