mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
Paweł Banaszewski
44eb71f922
Move the stateless AI Gateway router into `aibridge/x/proxy`. Extract shared routing helpers into `aibridge/routing` and move provider and recorder helpers to their owning packages. Update callers and tests without changing behavior. Generated by Coder Agents.
550 lines
19 KiB
Go
550 lines
19 KiB
Go
package aibridge_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.opentelemetry.io/otel"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"golang.org/x/xerrors"
|
|
|
|
"cdr.dev/slog/v3/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/aibridge"
|
|
"github.com/coder/coder/v2/aibridge/aibridgetest"
|
|
"github.com/coder/coder/v2/aibridge/config"
|
|
"github.com/coder/coder/v2/aibridge/intercept"
|
|
"github.com/coder/coder/v2/aibridge/internal/testutil"
|
|
"github.com/coder/coder/v2/aibridge/provider"
|
|
"github.com/coder/coder/v2/aibridge/routing"
|
|
"github.com/coder/coder/v2/coderd/httpapi"
|
|
codertestutil "github.com/coder/coder/v2/testutil"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
var bridgeTestTracer = otel.Tracer("bridge_test")
|
|
|
|
// TestRequestBridgeShutdownAdmissionRace deterministically interleaves request
|
|
// admission with Shutdown using the `serve_admission` quartz trap.
|
|
func TestRequestBridgeShutdownAdmissionRace(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
|
defer cancel()
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
|
|
|
release := make(chan struct{})
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
<-release
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
clk := quartz.NewMock(t)
|
|
trap := clk.Trap().Now("serve_admission")
|
|
defer trap.Close()
|
|
|
|
rec := testutil.MockRecorder{}
|
|
prov := aibridge.NewOpenAIProvider(config.OpenAI{BaseURL: upstream.URL})
|
|
bridge, err := aibridge.NewRequestBridge(ctx, []provider.Provider{prov}, &rec, nil, logger, nil, bridgeTestTracer, aibridge.WithClock(clk))
|
|
require.NoError(t, err)
|
|
|
|
serve := func(done chan struct{}) {
|
|
defer close(done)
|
|
bridge.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/openai/v1/conversations", nil))
|
|
}
|
|
|
|
// Request 1: admit past the trap; it then blocks in the upstream, holding
|
|
// the inflight WaitGroup (counter == 1).
|
|
req1 := make(chan struct{})
|
|
go serve(req1)
|
|
trap.MustWait(ctx).MustRelease(ctx)
|
|
|
|
// Request 2: park at the trap, having passed the closed check but before
|
|
// inflightWG.Add.
|
|
req2 := make(chan struct{})
|
|
go serve(req2)
|
|
call2 := trap.MustWait(ctx)
|
|
|
|
// Shutdown closes and waits on the inflight WaitGroup (held by request 1).
|
|
shutdown := make(chan struct{})
|
|
go func() {
|
|
defer close(shutdown)
|
|
_ = bridge.Shutdown(context.Background())
|
|
}()
|
|
|
|
// Releasing request 2 races its inflightWG.Add against Shutdown's Wait.
|
|
call2.MustRelease(ctx)
|
|
|
|
// Let both requests complete so Shutdown can finish.
|
|
close(release)
|
|
_ = codertestutil.TryReceive(ctx, t, req1)
|
|
_ = codertestutil.TryReceive(ctx, t, req2)
|
|
_ = codertestutil.TryReceive(ctx, t, shutdown)
|
|
|
|
// New requests are refused after graceful shutdown.
|
|
rejected := httptest.NewRecorder()
|
|
bridge.ServeHTTP(rejected, httptest.NewRequest(http.MethodGet, "/openai/v1/conversations", nil))
|
|
require.Equal(t, http.StatusServiceUnavailable, rejected.Code)
|
|
require.Equal(t, "AI Gateway is shutting down\n", rejected.Body.String())
|
|
}
|
|
|
|
// TestRequestBridgeShutdownDeadlineCancels verifies that deadline cancellation
|
|
// lets an in-flight proxy request finish without releasing the upstream manually.
|
|
func TestRequestBridgeShutdownDeadlineCancels(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
|
defer cancel()
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
|
|
|
// Set up an upstream that waits for cancellation or test cleanup.
|
|
started := make(chan struct{}, 1)
|
|
release := make(chan struct{})
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
started <- struct{}{}
|
|
select {
|
|
case <-release:
|
|
case <-r.Context().Done():
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
t.Cleanup(func() { close(release) })
|
|
|
|
// Send a request through the bridge to exercise upstream cancellation.
|
|
rec := testutil.MockRecorder{}
|
|
prov := aibridge.NewOpenAIProvider(config.OpenAI{BaseURL: upstream.URL})
|
|
bridge, err := aibridge.NewRequestBridge(ctx, []provider.Provider{prov}, &rec, nil, logger, nil, bridgeTestTracer)
|
|
require.NoError(t, err)
|
|
|
|
reqDone := make(chan struct{})
|
|
go func() {
|
|
defer close(reqDone)
|
|
bridge.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/openai/v1/conversations", nil))
|
|
}()
|
|
|
|
// Wait until the request is in flight in the upstream.
|
|
_ = codertestutil.TryReceive(ctx, t, started)
|
|
require.EqualValues(t, 1, bridge.InflightRequests())
|
|
|
|
// Shut down with an already-expired context to cancel the blocked request.
|
|
deadlineCtx, cancelDeadline := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
|
|
defer cancelDeadline()
|
|
require.ErrorIs(t, bridge.Shutdown(deadlineCtx), context.DeadlineExceeded)
|
|
|
|
// The request finishes after cancellation while release remains open.
|
|
// Completion need not precede Shutdown returning.
|
|
_ = codertestutil.TryReceive(ctx, t, reqDone)
|
|
require.Zero(t, bridge.InflightRequests())
|
|
}
|
|
|
|
func TestValidateProviders(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
logger := slogtest.Make(t, nil)
|
|
|
|
tests := []struct {
|
|
name string
|
|
providers []provider.Provider
|
|
expectErr string
|
|
}{
|
|
{
|
|
name: "all_supported_providers",
|
|
providers: []provider.Provider{
|
|
aibridge.NewOpenAIProvider(config.OpenAI{Name: "openai", BaseURL: "https://api.openai.com/v1/"}),
|
|
aibridgetest.NewAnthropicProvider(t, config.Anthropic{Name: "anthropic", BaseURL: "https://api.anthropic.com/"}, nil),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: "https://api.individual.githubcopilot.com"}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot-business", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot-enterprise", BaseURL: "https://api.enterprise.githubcopilot.com"}),
|
|
},
|
|
},
|
|
{
|
|
name: "default_names_and_base_urls",
|
|
providers: []provider.Provider{
|
|
aibridge.NewOpenAIProvider(config.OpenAI{}),
|
|
aibridgetest.NewAnthropicProvider(t, config.Anthropic{}, nil),
|
|
aibridge.NewCopilotProvider(config.Copilot{}),
|
|
},
|
|
},
|
|
{
|
|
name: "multiple_copilot_instances",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot-business", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot-enterprise", BaseURL: "https://api.enterprise.githubcopilot.com"}),
|
|
},
|
|
},
|
|
{
|
|
name: "name_with_slashes",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot/business", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
},
|
|
expectErr: "invalid provider name",
|
|
},
|
|
{
|
|
name: "name_with_spaces",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot business", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
},
|
|
expectErr: "invalid provider name",
|
|
},
|
|
{
|
|
name: "name_with_uppercase",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "Copilot", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
},
|
|
expectErr: "invalid provider name",
|
|
},
|
|
{
|
|
name: "unique_names",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: "https://api.individual.githubcopilot.com"}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot-business", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
},
|
|
},
|
|
{
|
|
name: "duplicate_base_url_different_names",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: "https://api.individual.githubcopilot.com"}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot-business", BaseURL: "https://api.individual.githubcopilot.com"}),
|
|
},
|
|
},
|
|
{
|
|
name: "duplicate_name",
|
|
providers: []provider.Provider{
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: "https://api.individual.githubcopilot.com"}),
|
|
aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: "https://api.business.githubcopilot.com"}),
|
|
},
|
|
expectErr: "duplicate provider name",
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
_, err := aibridge.NewRequestBridge(t.Context(), tc.providers, nil, nil, logger, nil, bridgeTestTracer)
|
|
if tc.expectErr != "" {
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), tc.expectErr)
|
|
} else {
|
|
require.NoError(t, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPassthroughRoutesForProviders(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
upstreamRespBody := "upstream response"
|
|
tests := []struct {
|
|
name string
|
|
baseURLPath string
|
|
requestMethod string
|
|
requestPath string
|
|
provider func(*testing.T, string) provider.Provider
|
|
expectPath string
|
|
}{
|
|
{
|
|
name: "openAI_no_base_path",
|
|
requestPath: "/openai/v1/conversations",
|
|
provider: func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewOpenAIProvider(config.OpenAI{BaseURL: baseURL})
|
|
},
|
|
expectPath: "/conversations",
|
|
},
|
|
{
|
|
name: "openAI_with_base_path",
|
|
baseURLPath: "/v1",
|
|
requestPath: "/openai/v1/conversations",
|
|
provider: func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewOpenAIProvider(config.OpenAI{BaseURL: baseURL})
|
|
},
|
|
expectPath: "/v1/conversations",
|
|
},
|
|
{
|
|
name: "anthropic_no_base_path",
|
|
requestPath: "/anthropic/v1/models",
|
|
provider: func(t *testing.T, baseURL string) provider.Provider {
|
|
return aibridgetest.NewAnthropicProvider(t, config.Anthropic{BaseURL: baseURL}, nil)
|
|
},
|
|
expectPath: "/v1/models",
|
|
},
|
|
{
|
|
name: "anthropic_with_base_path",
|
|
baseURLPath: "/v1",
|
|
requestPath: "/anthropic/v1/models",
|
|
provider: func(t *testing.T, baseURL string) provider.Provider {
|
|
return aibridgetest.NewAnthropicProvider(t, config.Anthropic{BaseURL: baseURL}, nil)
|
|
},
|
|
expectPath: "/v1/v1/models",
|
|
},
|
|
{
|
|
name: "copilot_no_base_path",
|
|
requestPath: "/copilot/models",
|
|
provider: func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL})
|
|
},
|
|
expectPath: "/models",
|
|
},
|
|
{
|
|
name: "copilot_with_base_path",
|
|
baseURLPath: "/v1",
|
|
requestPath: "/copilot/models",
|
|
provider: func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL})
|
|
},
|
|
expectPath: "/v1/models",
|
|
},
|
|
{
|
|
name: "copilot_unknown_route",
|
|
requestMethod: http.MethodPost,
|
|
requestPath: "/copilot/future/endpoint",
|
|
provider: func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewCopilotProvider(config.Copilot{BaseURL: baseURL})
|
|
},
|
|
expectPath: "/future/endpoint",
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
logger := slogtest.Make(t, nil)
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
assert.Equal(t, tc.expectPath, r.URL.Path)
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(upstreamRespBody))
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
rec := testutil.MockRecorder{}
|
|
prov := tc.provider(t, upstream.URL+tc.baseURLPath)
|
|
bridge, err := aibridge.NewRequestBridge(t.Context(), []provider.Provider{prov}, &rec, nil, logger, nil, bridgeTestTracer)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(tc.requestMethod, tc.requestPath, nil)
|
|
resp := httptest.NewRecorder()
|
|
bridge.ServeHTTP(resp, req)
|
|
|
|
assert.Equal(t, http.StatusOK, resp.Code)
|
|
assert.Contains(t, resp.Body.String(), upstreamRespBody)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBridgedRouteTakesPrecedenceOverPassthroughCatchAll(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
upstreamCalled := false
|
|
interceptorCalled := false
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
upstreamCalled = true
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
// Given: a provider with a bridged route and a passthrough catch-all.
|
|
prov := &testutil.MockProvider{
|
|
NameStr: "test",
|
|
URL: upstream.URL,
|
|
Bridged: []string{"/responses"},
|
|
Passthrough: []string{"/"},
|
|
InterceptorFunc: func(http.ResponseWriter, *http.Request, trace.Tracer) (intercept.Interceptor, error) {
|
|
interceptorCalled = true
|
|
return nil, xerrors.New("test interceptor error")
|
|
},
|
|
}
|
|
bridge, err := aibridge.NewRequestBridge(
|
|
t.Context(),
|
|
[]provider.Provider{prov},
|
|
nil, nil, slogtest.Make(t, nil), nil, bridgeTestTracer,
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
// When: a request targets the bridged route.
|
|
req := httptest.NewRequest(http.MethodPost, "/test/responses", nil)
|
|
resp := httptest.NewRecorder()
|
|
bridge.ServeHTTP(resp, req)
|
|
|
|
// Then: the interceptor handles it, not the passthrough upstream.
|
|
assert.Equal(t, http.StatusInternalServerError, resp.Code)
|
|
assert.True(t, interceptorCalled)
|
|
assert.False(t, upstreamCalled)
|
|
}
|
|
|
|
func TestWebSocketUpgradeRejected(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
interceptorCalled := false
|
|
prov := &testutil.MockProvider{
|
|
NameStr: "test",
|
|
Bridged: []string{"/responses"},
|
|
InterceptorFunc: func(http.ResponseWriter, *http.Request, trace.Tracer) (intercept.Interceptor, error) {
|
|
interceptorCalled = true
|
|
return nil, nil //nolint:nilnil // The interceptor must not be reached.
|
|
},
|
|
}
|
|
bridge, err := aibridge.NewRequestBridge(
|
|
t.Context(),
|
|
[]provider.Provider{prov},
|
|
nil, nil, slogtest.Make(t, nil), nil, bridgeTestTracer,
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/test/responses", nil)
|
|
req.Header.Set("Connection", "keep-alive, Upgrade")
|
|
req.Header.Set("Upgrade", "WebSocket")
|
|
resp := httptest.NewRecorder()
|
|
|
|
bridge.ServeHTTP(resp, req)
|
|
|
|
assert.Equal(t, http.StatusNotImplemented, resp.Code)
|
|
assert.Contains(t, resp.Body.String(), "WebSocket transport is not supported, use HTTP")
|
|
assert.False(t, interceptorCalled)
|
|
}
|
|
|
|
func TestRequestBodySizeLimit(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
newOpenAI := func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewOpenAIProvider(config.OpenAI{Name: "openai", BaseURL: baseURL})
|
|
}
|
|
newAnthropic := func(t *testing.T, baseURL string) provider.Provider {
|
|
return aibridgetest.NewAnthropicProvider(t, config.Anthropic{Name: "anthropic", BaseURL: baseURL}, nil)
|
|
}
|
|
newCopilot := func(_ *testing.T, baseURL string) provider.Provider {
|
|
return aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: baseURL})
|
|
}
|
|
|
|
// Each body is a well-formed, schema-valid request for its provider, with
|
|
// an oversized message content that pushes it past the 32 MiB limit.
|
|
filler := strings.Repeat("A", 32<<20)
|
|
chatCompletionsBody := fmt.Appendf(nil, `{"model":"gpt-4","messages":[{"role":"user","content":"%s"}]}`, filler)
|
|
responsesBody := fmt.Appendf(nil, `{"model":"gpt-4","input":"%s"}`, filler)
|
|
messagesBody := fmt.Appendf(nil, `{"model":"claude-3-5-sonnet-latest","max_tokens":1024,"messages":[{"role":"user","content":"%s"}]}`, filler)
|
|
|
|
tests := []struct {
|
|
name string
|
|
provider func(*testing.T, string) provider.Provider
|
|
path string
|
|
body []byte
|
|
}{
|
|
{name: "openai_passthrough", provider: newOpenAI, path: "/openai/v1/models", body: chatCompletionsBody},
|
|
{name: "openai_chat_completions", provider: newOpenAI, path: "/openai/v1/chat/completions", body: chatCompletionsBody},
|
|
{name: "openai_responses", provider: newOpenAI, path: "/openai/v1/responses", body: responsesBody},
|
|
{name: "anthropic_passthrough", provider: newAnthropic, path: "/anthropic/v1/models", body: messagesBody},
|
|
{name: "anthropic_messages", provider: newAnthropic, path: "/anthropic/v1/messages", body: messagesBody},
|
|
{name: "copilot_passthrough", provider: newCopilot, path: "/copilot/models", body: chatCompletionsBody},
|
|
{name: "copilot_chat_completions", provider: newCopilot, path: "/copilot/chat/completions", body: chatCompletionsBody},
|
|
{name: "copilot_responses", provider: newCopilot, path: "/copilot/responses", body: responsesBody},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
logger := slogtest.Make(t, nil)
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = io.ReadAll(r.Body)
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
prov := tc.provider(t, upstream.URL)
|
|
bridge, err := aibridge.NewRequestBridge(
|
|
t.Context(),
|
|
[]provider.Provider{prov},
|
|
nil, nil, logger, nil, bridgeTestTracer,
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, tc.path, bytes.NewReader(tc.body))
|
|
// Unknown Content-Length
|
|
req.ContentLength = -1
|
|
// Copilot's bridged route checks Authorization before reading the
|
|
// body, so provide a token to reach the read path.
|
|
req.Header.Set("Authorization", "Bearer test-key")
|
|
// coderd mounts aibridge behind the middleware that turns this
|
|
// tracker into the too-large metric's reason label, so a rejection
|
|
// that leaves it unset is counted as some other kind of 413.
|
|
var tracker httpapi.RequestBodyLimitTracker
|
|
req = req.WithContext(httpapi.WithRequestBodyLimitTracker(req.Context(), &tracker))
|
|
resp := httptest.NewRecorder()
|
|
bridge.ServeHTTP(resp, req)
|
|
|
|
assert.Equal(t, http.StatusRequestEntityTooLarge, resp.Code)
|
|
assert.Contains(t, resp.Body.String(), "Request body too large")
|
|
assert.True(t, tracker.Exceeded(), "rejection must be attributed to the body size limit")
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestDisabledProviderHandler asserts that requests to a disabled
|
|
// provider return a 503 with an ErrorCodeProviderDisabled body and
|
|
// that a sibling enabled provider keeps routing normally.
|
|
func TestDisabledProviderHandler(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
logger := slogtest.Make(t, nil)
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
_, _ = w.Write([]byte("upstream-reached"))
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
enabled := aibridge.NewOpenAIProvider(config.OpenAI{Name: "enabled-openai", BaseURL: upstream.URL})
|
|
disabled := aibridge.NewDisabledProviderStub("disabled-openai", "openai")
|
|
bridge, err := aibridge.NewRequestBridge(
|
|
t.Context(),
|
|
[]provider.Provider{enabled, disabled},
|
|
nil, nil, logger, nil, bridgeTestTracer,
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
path string
|
|
}{
|
|
{name: "Bridged", path: "/disabled-openai/v1/chat/completions"},
|
|
{name: "Passthrough", path: "/disabled-openai/v1/models"},
|
|
{name: "Unknown", path: "/disabled-openai/anything/else"},
|
|
} {
|
|
t.Run("DisabledProviderReturnsSentinel/"+tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
req := httptest.NewRequest(http.MethodPost, tc.path, nil)
|
|
resp := httptest.NewRecorder()
|
|
bridge.ServeHTTP(resp, req)
|
|
|
|
assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
|
|
assert.Contains(t, resp.Body.String(), routing.ErrorCodeProviderDisabled)
|
|
assert.Contains(t, resp.Body.String(), "disabled-openai")
|
|
})
|
|
}
|
|
|
|
t.Run("EnabledProviderUnaffected", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/enabled-openai/v1/models", nil)
|
|
resp := httptest.NewRecorder()
|
|
bridge.ServeHTTP(resp, req)
|
|
|
|
assert.Equal(t, http.StatusOK, resp.Code)
|
|
assert.Equal(t, "upstream-reached", resp.Body.String())
|
|
})
|
|
}
|