llmgw/internal/control/external_test.go

package control

import (
	"context"
	"log/slog"
	"os"
	"path/filepath"
	"sync"
	"testing"

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

type recordSink struct {
	mu    sync.Mutex
	infos []endpoint.Info
	times int
}

func (s *recordSink) SetEndpoints(_ context.Context, infos []endpoint.Info) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.infos, s.times = infos, s.times+1
}

type noDispatch struct{}

func (noDispatch) StartShard(context.Context, reconciler.ShardCommand) error { return nil }
func (noDispatch) StopShard(context.Context, string, string) error           { return nil }
func (noDispatch) DrainShard(context.Context, string, string) error          { return nil }

func TestAnExternalModelIsPublishedWithItsKeyAndNothingIsPlacedForIt(t *testing.T) {
	dir := t.TempDir()
	keyFile := filepath.Join(dir, "provider.key")
	if err := os.WriteFile(keyFile, []byte("sk-provider-1\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	models := []catalog.Model{
		{Name: "cloud", Source: "external", Aliases: []string{"gpt"}, Scaling: catalog.Scaling{Min: 3, Max: 3},
			External: &catalog.External{URL: "https://api.example.com/", Model: "gpt-4o-mini", APIKeyFile: keyFile}},
		{Name: "gone", Source: "external",
			External: &catalog.External{URL: "https://api.example.com", Model: "m", APIKeyFile: filepath.Join(dir, "missing")}},
	}
	cat, err := catalog.New(models)
	if err != nil {
		t.Fatal(err)
	}
	sink := &recordSink{}
	c := New(DefaultConfig(), catalog.NewLive(cat), noDispatch{}, sink, slog.New(slog.DiscardHandler))
	c.Publish(t.Context())

	if len(sink.infos) != 1 {
		t.Fatalf("published %+v; the model whose key is missing must not be offered", sink.infos)
	}
	e := sink.infos[0]
	if e.ID != "external:cloud" || e.URL != "https://api.example.com" || e.APIKey != "sk-provider-1" ||
		e.UpstreamModel != "gpt-4o-mini" || e.Engine != endpoint.EngineExternal || len(e.Models) != 2 {
		t.Errorf("endpoint %+v", e)
	}
	// A key rotated on disk is published on the next pass.
	_ = os.WriteFile(keyFile, []byte("sk-provider-2"), 0o600)
	c.Publish(t.Context())
	if sink.times != 2 || sink.infos[0].APIKey != "sk-provider-2" {
		t.Errorf("the rotated key was not published: %d pushes, key %q", sink.times, sink.infos[0].APIKey)
	}
	// Scaling of 3 or not, nothing is wanted of the cluster for a model
	// served elsewhere, and so it never reads as short or unplaceable.
	for _, h := range c.Reconciler().ModelHealth() {
		if h.Model == "cloud" && (h.Desired != 0 || h.Unplaceable) {
			t.Errorf("an external model wants replicas: %+v", h)
		}
	}
}

// Several keys for one provider: an endpoint each, numbered, with its
// weight; one whose file cannot be read is left out and the others stay.
func TestEachOfAProvidersKeysIsAnEndpointOfItsOwn(t *testing.T) {
	dir := t.TempDir()
	a, c := filepath.Join(dir, "a.key"), filepath.Join(dir, "c.key")
	_ = os.WriteFile(a, []byte("sk-a\n"), 0o600)
	_ = os.WriteFile(c, []byte("sk-c"), 0o600)
	cat, err := catalog.New([]catalog.Model{{Name: "cloud", Source: "external",
		External: &catalog.External{URL: "https://api.example.com", Model: "gpt-4o-mini",
			APIKeyFiles: []catalog.ExternalKey{{File: a, Weight: 3}, {File: filepath.Join(dir, "gone")}, {File: c}}}}})
	if err != nil {
		t.Fatal(err)
	}
	sink := &recordSink{}
	New(DefaultConfig(), catalog.NewLive(cat), noDispatch{}, sink, slog.New(slog.DiscardHandler)).Publish(t.Context())
	got := map[string]endpoint.Info{}
	for _, e := range sink.infos {
		got[e.ID] = e
	}
	if len(got) != 2 || got["external:cloud#1"].APIKey != "sk-a" || got["external:cloud#1"].Weight != 3 ||
		got["external:cloud#3"].APIKey != "sk-c" || got["external:cloud#3"].Weight != 1 {
		t.Errorf("published %+v", sink.infos)
	}
}