Function TestClassifyModelError

pkg/modelerrors/modelerrors_test.go:776–846

(t *testing.T)

}

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

	tests := []struct {
		name            string
		err             error
		wantRetryable   bool
		wantRateLimited bool
		wantRetryAfter  time.Duration
	}{
		{name: "nil", err: nil, wantRetryable: false, wantRateLimited: false},
		{name: "context canceled", err: context.Canceled, wantRetryable: false, wantRateLimited: false},
		{name: "context deadline exceeded", err: context.DeadlineExceeded, wantRetryable: false, wantRateLimited: false},
		{name: "context overflow", err: errors.New("prompt is too long: 200000 tokens > 100000 maximum"), wantRetryable: false, wantRateLimited: false},
		// 429 without StatusError (fallback message-pattern path)
		{name: "429 message fallback, no RetryAfter", err: errors.New("POST /v1/chat: 429 Too Many Requests"), wantRetryable: false, wantRateLimited: true, wantRetryAfter: 0},
		// 429 via StatusError (primary path) — no Retry-After
		{name: "429 StatusError no retry-after", err: &StatusError{StatusCode: 429, RetryAfter: 0, Err: errors.New("rate limited")}, wantRetryable: false, wantRateLimited: true, wantRetryAfter: 0},
		// 429 via StatusError with Retry-After from response header
		{name: "429 StatusError with retry-after", err: &StatusError{StatusCode: 429, RetryAfter: 20 * time.Second, Err: errors.New("rate limited")}, wantRetryable: false, wantRateLimited: true, wantRetryAfter: 20 * time.Second},
		// Retryable status codes via StatusError
		{name: "500 StatusError", err: &StatusError{StatusCode: 500, Err: errors.New("internal server error")}, wantRetryable: true, wantRateLimited: false},
		{name: "529 StatusError", err: &StatusError{StatusCode: 529, Err: errors.New("overloaded")}, wantRetryable: true, wantRateLimited: false},
		{name: "408 StatusError", err: &StatusError{StatusCode: 408, Err: errors.New("timeout")}, wantRetryable: true, wantRateLimited: false},
		// Retryable fallback path (message-based)
		{name: "500 message fallback", err: errors.New("500 internal server error"), wantRetryable: true, wantRateLimited: false},
		{name: "502 message fallback", err: errors.New("502 bad gateway"), wantRetryable: true, wantRateLimited: false},
		// Non-retryable via StatusError
		{name: "401 StatusError", err: &StatusError{StatusCode: 401, Err: errors.New("unauthorized")}, wantRetryable: false, wantRateLimited: false},
		{name: "403 StatusError", err: &StatusError{StatusCode: 403, Err: errors.New("forbidden")}, wantRetryable: false, wantRateLimited: false},
		// Non-retryable fallback
		{name: "401 message fallback", err: errors.New("401 unauthorized"), wantRetryable: false, wantRateLimited: false},
		// 400 with Vertex AI "function response parts" message is treated as transient (issue #2683)
		{name: "vertex transient 400 StatusError", err: &StatusError{StatusCode: 400, Err: errors.New("Error 400, Message: Please ensure that the number of function response parts is equal to the number of function call parts of the function call turn., Status: INVALID_ARGUMENT, Details: []")}, wantRetryable: true, wantRateLimited: false},
		{name: "vertex transient 400 wrapped in stream error", err: fmt.Errorf("error receiving from stream: %w", &StatusError{StatusCode: 400, Err: errors.New("number of function response parts")}), wantRetryable: true, wantRateLimited: false},
		{name: "vertex transient 400 message fallback (no StatusError)", err: errors.New("400 Bad Request: Please ensure that the number of function response parts is equal to the number of function call parts"), wantRetryable: true, wantRateLimited: false},
		// Network errors
		{name: "network timeout", err: &mockTimeoutError{}, wantRetryable: true, wantRateLimited: false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			retryable, rateLimited, retryAfterOut := ClassifyModelError(tt.err)
			assert.Equal(t, tt.wantRetryable, retryable, "retryable mismatch")
			assert.Equal(t, tt.wantRateLimited, rateLimited, "rateLimited mismatch")
			assert.Equal(t, tt.wantRetryAfter, retryAfterOut, "retryAfter mismatch")
		})
	}

	t.Run("wrapped StatusError is found by errors.As", func(t *testing.T) {
		t.Parallel()
		statusErr := &StatusError{StatusCode: 429, RetryAfter: 15 * time.Second, Err: errors.New("rate limited")}
		wrapped := fmt.Errorf("model failed: %w", statusErr)
		retryable, rateLimited, retryAfterOut := ClassifyModelError(wrapped)
		assert.False(t, retryable)
		assert.True(t, rateLimited)
		assert.Equal(t, 15*time.Second, retryAfterOut)