internal/gateway/request_model_resolution_test.go
Rama main · réplica de lectura del repositorio original.
package gateway
import (
"context"
"errors"
"io"
"net/http"
"testing"
"github.com/enterpilot/gomodel/internal/core"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type requestRefreshProvider struct {
supported map[string]bool
providerType map[string]string
modelCount int
refreshErr error
resolveErrWhenEmpty bool
refreshCalls int
resolveErr error
}
func newRequestRefreshProvider(modelCount int) *requestRefreshProvider {
return &requestRefreshProvider{
supported: map[string]bool{
"openai/gpt-4o": true,
},
providerType: map[string]string{
"openai/gpt-4o": "openai",
},
modelCount: modelCount,
}
}
func (p *requestRefreshProvider) RefreshProviderModels(_ context.Context, providerSelector string) (int, error) {
p.refreshCalls++
if p.refreshErr != nil {
return 0, p.refreshErr
}
if providerSelector != "ollama" {
return 0, nil
}
p.supported["ollama/qwen3:8b"] = true
p.providerType["ollama/qwen3:8b"] = "ollama"
p.modelCount = 1
return 1, nil
}
func (p *requestRefreshProvider) ResolveModel(requested core.RequestedModelSelector) (core.ModelSelector, bool, error) {
if p.resolveErr != nil {
return core.ModelSelector{}, false, p.resolveErr
}
if p.resolveErrWhenEmpty && p.modelCount == 0 {
return core.ModelSelector{}, false, core.NewProviderError("", http.StatusServiceUnavailable, "model registry not initialized", nil)
}
selector, err := requested.Normalize()
return selector, false, err
}
func (p *requestRefreshProvider) Supports(model string) bool {
return p.supported[model]
}
func (p *requestRefreshProvider) GetProviderType(model string) string {
return p.providerType[model]
}
func (p *requestRefreshProvider) ModelCount() int {
return p.modelCount
}
func (p *requestRefreshProvider) ChatCompletion(context.Context, *core.ChatRequest) (*core.ChatResponse, error) {
return nil, nil
}
func (p *requestRefreshProvider) StreamChatCompletion(context.Context, *core.ChatRequest) (io.ReadCloser, error) {
return nil, nil
}
func (p *requestRefreshProvider) ListModels(context.Context) (*core.ModelsResponse, error) {
return nil, nil
}
func (p *requestRefreshProvider) Responses(context.Context, *core.ResponsesRequest) (*core.ResponsesResponse, error) {
return nil, nil
}
func (p *requestRefreshProvider) StreamResponses(context.Context, *core.ResponsesRequest) (io.ReadCloser, error) {
return nil, nil
}
func (p *requestRefreshProvider) Embeddings(context.Context, *core.EmbeddingRequest) (*core.EmbeddingResponse, error) {
return nil, nil
}
type requestAliasResolver map[string]core.ModelSelector
func (r requestAliasResolver) ResolveModel(requested core.RequestedModelSelector) (core.ModelSelector, bool, error) {
if selector, ok := r[requested.RequestedQualifiedModel()]; ok {
return selector, true, nil
}
selector, err := requested.Normalize()
return selector, false, err
}
type requestSlowdownResolver struct {
requestAliasResolver
factor float64
requested string
resolved string
}
func (r *requestSlowdownResolver) ResolveSlowdown(_ context.Context, requested core.RequestedModelSelector, resolved core.ModelSelector) float64 {
r.requested = requested.RequestedQualifiedModel()
r.resolved = resolved.QualifiedModel()
return r.factor
}
func TestResolveRequestModelCarriesResolvedSlowdownFactor(t *testing.T) {
aliasResolver := &requestSlowdownResolver{
requestAliasResolver: requestAliasResolver{
"smart": {Provider: "openai", Model: "gpt-4o"},
},
factor: 0.5,
}
directResolver := &requestSlowdownResolver{factor: 0.3}
tests := []struct {
name string
resolver ModelResolver
requested string
wantSlowdown float64
capture *requestSlowdownResolver
wantResolved string
}{
{name: "alias factor", resolver: aliasResolver, requested: "smart", wantSlowdown: 0.5, capture: aliasResolver, wantResolved: "openai/gpt-4o"},
{name: "direct model factor", resolver: directResolver, requested: "openai/gpt-4o", wantSlowdown: 0.3, capture: directResolver, wantResolved: "openai/gpt-4o"},
{name: "resolver without slowdown contract", resolver: requestAliasResolver{}, requested: "openai/gpt-4o", wantSlowdown: 0, wantResolved: "openai/gpt-4o"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resolution, err := ResolveRequestModelWithAuthorizer(
context.Background(), newRequestRefreshProvider(1), tt.resolver, nil,
core.NewRequestedModelSelector(tt.requested, ""),
)
require.NoError(t, err)
require.Equal(t, tt.wantSlowdown, resolution.Slowdown)
got := resolution.ResolvedSelector.QualifiedModel()
require.Equal(t, tt.wantResolved, got)
if tt.capture != nil {
require.Equal(t, tt.requested, tt.capture.requested, "slowdown resolver requested model")
require.Equal(t, tt.wantResolved, tt.capture.resolved, "slowdown resolver resolved model")
}
})
}
}
type requestRefreshTargetResolver struct {
provider *requestRefreshProvider
target core.ModelSelector
err error
}
func (r requestRefreshTargetResolver) ResolveModel(requested core.RequestedModelSelector) (core.ModelSelector, bool, error) {
if requested.RequestedQualifiedModel() == "smart" && r.provider.Supports(r.target.QualifiedModel()) {
return r.target, true, nil
}
selector, err := requested.Normalize()
return selector, false, err
}
func (r requestRefreshTargetResolver) ResolveRefreshTarget(requested core.RequestedModelSelector) (core.ModelSelector, bool, error) {
if requested.RequestedQualifiedModel() != "smart" {
return core.ModelSelector{}, false, nil
}
if r.err != nil {
return core.ModelSelector{}, false, r.err
}
return r.target, true, nil
}
func TestResolveRequestModelRefreshesBeforeUnsupportedModel(t *testing.T) {
provider := newRequestRefreshProvider(1)
resolution, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
nil,
nil,
core.NewRequestedModelSelector("ollama/qwen3:8b", ""),
)
require.NoError(t, err)
require.Equal(t, 1, provider.refreshCalls)
got := resolution.ResolvedQualifiedModel()
require.Equal(t, "ollama/qwen3:8b", got)
got = resolution.ProviderType
require.Equal(t, "ollama", got)
}
func TestResolveRequestModelRefreshesBeforeEmptyRegistryFailure(t *testing.T) {
provider := newRequestRefreshProvider(0)
resolution, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
nil,
nil,
core.NewRequestedModelSelector("ollama/qwen3:8b", ""),
)
require.NoError(t, err)
require.Equal(t, 1, provider.refreshCalls)
got := resolution.ResolvedQualifiedModel()
require.Equal(t, "ollama/qwen3:8b", got)
}
func TestResolveRequestModelRefreshesAliasTargetBeforeCatalogSupportsIt(t *testing.T) {
provider := newRequestRefreshProvider(1)
resolver := requestRefreshTargetResolver{
provider: provider,
target: core.ModelSelector{Provider: "ollama", Model: "qwen3:8b"},
}
resolution, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
resolver,
nil,
core.NewRequestedModelSelector("smart", ""),
)
require.NoError(t, err)
require.Equal(t, 1, provider.refreshCalls)
got := resolution.ResolvedQualifiedModel()
require.Equal(t, "ollama/qwen3:8b", got)
require.True(t, resolution.AliasApplied)
}
func TestResolveRequestModelReturnsRefreshTargetError(t *testing.T) {
provider := newRequestRefreshProvider(1)
targetErr := errors.New("invalid alias target")
resolver := requestRefreshTargetResolver{
provider: provider,
target: core.ModelSelector{Provider: "ollama", Model: "qwen3:8b"},
err: targetErr,
}
_, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
resolver,
nil,
core.NewRequestedModelSelector("smart", ""),
)
require.Error(t, err)
require.ErrorIs(t, err, targetErr)
require.Equal(t, 0, provider.refreshCalls)
}
func TestResolveRequestModelRefreshesAliasTargetAfterResolverFailure(t *testing.T) {
provider := newRequestRefreshProvider(0)
provider.resolveErrWhenEmpty = true
resolver := requestAliasResolver{
"smart": {Provider: "ollama", Model: "qwen3:8b"},
}
resolution, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
resolver,
nil,
core.NewRequestedModelSelector("smart", ""),
)
require.NoError(t, err)
require.Equal(t, 1, provider.refreshCalls)
got := resolution.ResolvedQualifiedModel()
require.Equal(t, "ollama/qwen3:8b", got)
require.True(t, resolution.AliasApplied)
}
func TestResolveRequestModelReturnsRefreshError(t *testing.T) {
provider := newRequestRefreshProvider(1)
provider.refreshErr = core.NewProviderError("ollama", http.StatusServiceUnavailable, "provider is unavailable", errors.New("connection refused"))
_, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
nil,
nil,
core.NewRequestedModelSelector("ollama/qwen3:8b", ""),
)
require.Error(t, err)
var gatewayErr *core.GatewayError
require.ErrorAs(t, err, &gatewayErr)
require.Equal(t, http.StatusServiceUnavailable, gatewayErr.HTTPStatusCode())
require.Equal(t, core.ErrorTypeProvider, gatewayErr.Type)
}
// nonRefreshingProvider hides RefreshProviderModels, for a provider whose
// models cannot be refreshed at request time.
type nonRefreshingProvider struct {
core.RoutableProvider
inner *requestRefreshProvider
}
func (p nonRefreshingProvider) ResolveModel(requested core.RequestedModelSelector) (core.ModelSelector, bool, error) {
return p.inner.ResolveModel(requested)
}
func (p nonRefreshingProvider) ModelCount() int { return p.inner.ModelCount() }
// A resolver's typed gateway error keeps its status: while the model registry
// is still empty the router answers 503 so clients retry, and that must not be
// rewritten into a 400 that blames the request.
func TestResolveRequestModelKeepsResolverErrorStatus(t *testing.T) {
tests := []struct {
name string
modelCount int
resolveEmpty bool // the resolver itself reports the empty registry
resolveErr error // a plain resolver error
noRefresh bool
wantStatus int
wantType core.ErrorType
}{
{name: "registry not ready", resolveEmpty: true,
wantStatus: http.StatusServiceUnavailable, wantType: core.ErrorTypeProvider},
{name: "registry not ready without refresh", resolveEmpty: true, noRefresh: true,
wantStatus: http.StatusServiceUnavailable, wantType: core.ErrorTypeProvider},
{name: "empty registry with a resolving selector",
wantStatus: http.StatusServiceUnavailable, wantType: core.ErrorTypeProvider},
{name: "plain resolver error", modelCount: 1, resolveErr: errors.New("unknown alias"),
wantStatus: http.StatusBadRequest, wantType: core.ErrorTypeInvalidRequest},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
inner := newRequestRefreshProvider(tt.modelCount)
inner.resolveErrWhenEmpty = tt.resolveEmpty
inner.resolveErr = tt.resolveErr
var provider core.RoutableProvider = inner
if tt.noRefresh {
provider = nonRefreshingProvider{RoutableProvider: inner, inner: inner}
}
_, err := ResolveRequestModelWithAuthorizer(
context.Background(),
provider,
nil,
nil,
core.NewRequestedModelSelector("openai/gpt-4o", ""),
)
var gatewayErr *core.GatewayError
require.ErrorAs(t, err, &gatewayErr)
assert.Equal(t, tt.wantStatus, gatewayErr.HTTPStatusCode())
assert.Equal(t, tt.wantType, gatewayErr.Type)
})
}
}
// requestRouteResolver reports the virtual models its route passed through.
type requestRouteResolver struct {
requestAliasResolver
path []string
}
func (r requestRouteResolver) ResolveModelRouteForUserPath(_ context.Context, requested core.RequestedModelSelector) (core.ModelSelector, []string, bool, error) {
selector, changed, err := r.ResolveModel(requested)
return selector, r.path, changed, err
}
func TestResolveRequestModelCarriesVirtualModelPath(t *testing.T) {
resolver := requestRouteResolver{
requestAliasResolver: requestAliasResolver{"team": {Provider: "openai", Model: "gpt-4o"}},
path: []string{"team", "smart"},
}
resolution, err := ResolveRequestModelWithAuthorizer(
context.Background(), newRequestRefreshProvider(1), resolver, nil,
core.NewRequestedModelSelector("team", ""),
)
require.NoError(t, err)
assert.Equal(t, []string{"team", "smart"}, resolution.VirtualModelPath)
}