Function TestPassthrough_KeyFailover

aibridge/passthrough_internal_test.go:298–594

TestPassthrough_KeyFailover exercises the KeyFailoverTransport end-to-end through the passthrough proxy, parameterised over providers (anthropic, openai). Each scenario asserts the upstream request count, the response status and Retry-After, and the final pool state.

(t *testing.T)

// request count, the response status and Retry-After, and the final
// pool state.
func TestPassthrough_KeyFailover(t *testing.T) {
	t.Parallel()

	type upstreamResponse struct {
		statusCode int
		body       string
		headers    map[string]string
	}

	const (
		rateLimitBody   = `{"error":"rate"}`
		authErrorBody   = `{"error":"unauthorized"}`
		serverErrorBody = `{"error":"server"}`
		successBody     = `{"data":[]}`
	)

	// providers parameterises the table over the providers exposed
	// to the failover transport. Each entry encapsulates the
	// provider-specific bits the test needs: how the mock upstream
	// extracts the key from the request, how a BYOK request sets
	// it, and how the provider is constructed for a given pool.
	providers := []struct {
		name        string
		byokOnly    bool
		extractKey  func(*http.Request) string
		setBYOK     func(*http.Request, string)
		newProvider func(baseURL string, pool *keypool.Pool) provider.Provider
	}{
		{
			name: "anthropic",
			extractKey: func(r *http.Request) string {
				return r.Header.Get("X-Api-Key")
			},
			setBYOK: func(r *http.Request, key string) {
				r.Header.Set("X-Api-Key", key)
			},
			newProvider: func(baseURL string, pool *keypool.Pool) provider.Provider {
				return provider.NewAnthropic(config.Anthropic{
					BaseURL: baseURL,
					KeyPool: pool,
				}, nil)
			},
		},
		{
			name: "openai",
			extractKey: func(r *http.Request) string {
				return strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
			},
			setBYOK: func(r *http.Request, key string) {
				r.Header.Set("Authorization", "Bearer "+key)
			},
			newProvider: func(baseURL string, pool *keypool.Pool) provider.Provider {
				cfg := config.OpenAI{BaseURL: baseURL}
				if pool != nil {
					cfg.KeyPool = pool
				}
				return provider.NewOpenAI(cfg)
			},