llmgw/internal/router/residency_test.go

package router

import (
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

func TestAProjectInTheClusterIsNeverSentOutside(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	var outside atomic.Int64
	provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		outside.Add(1)
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"from outside"}}]}`))
	}))
	t.Cleanup(provider.Close)
	local := catModel("qwen2.5-14b-instruct", []string{"q14"}, "cloud")
	cloud := catalog.Model{Name: "cloud", Source: "external", External: &catalog.External{URL: provider.URL, Model: "m"}}
	cat, err := catalog.New([]catalog.Model{local, cloud})
	if err != nil {
		t.Fatal(err)
	}
	rt.SetCatalog(cat)
	rt.SetEndpoints(t.Context(), []endpoint.Info{
		{ID: "a", URL: e.srv.URL, Models: []string{"qwen2.5-14b-instruct"}},
		{ID: "external:cloud", URL: provider.URL, Models: []string{"cloud"}, Engine: endpoint.EngineExternal,
			Pool: "external", Weight: 1, UpstreamModel: "m"},
	})
	auth := map[string]string{"Authorization": "Bearer " + secret}

	if err := rt.store.SetInCluster(t.Context(), mustProjectOf(t, rt, "key-1"), true); err != nil {
		t.Fatal(err)
	}
	if rec := post(t, rt.Handler(), map[string]any{"model": "cloud", "messages": []any{}}, auth); rec.Code != http.StatusForbidden {
		t.Errorf("an outside model, asked for: %d %s", rec.Code, rec.Body.String())
	}
	// The local model down: its outside fallback is not used.
	e.status.Store(http.StatusServiceUnavailable)
	if rec := post(t, rt.Handler(), map[string]any{"model": "q14", "messages": []any{}}, auth); rec.Code == http.StatusOK {
		t.Errorf("served: %s", rec.Body.String())
	}
	if outside.Load() != 0 {
		t.Errorf("the outside provider was sent %d requests", outside.Load())
	}
}