llmgw/internal/router/chat_load_test.go

package router

import (
	"encoding/json"
	"net/http/httptest"
	"strings"
	"testing"

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

// replicaSource is the control plane's view with only replicas in it.
type replicaSource struct {
	replicas map[string][]adminapi.ReplicaView
}

func (s replicaSource) Nodes() []adminapi.NodeView   { return nil }
func (s replicaSource) Models() []adminapi.ModelView { return nil }
func (s replicaSource) Replicas(model string) []adminapi.ReplicaView {
	return s.replicas[model]
}
func (s replicaSource) Explain(string) adminapi.ExplainView { return adminapi.ExplainView{} }
func (s replicaSource) Summary() adminapi.Summary           { return adminapi.Summary{} }

func onDemandRouter(t *testing.T) (*Router, *fakeActivator) {
	t.Helper()
	e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	rt, _ := chatRouter(t, e.srv.URL)
	cpu := []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}
	c, err := catalog.New([]catalog.Model{
		{Name: "qwen2.5-coder-7b-instruct", Source: "http://m/a.gguf", Profiles: cpu},
		{Name: "nemo", Source: "http://m/b.gguf", Profiles: cpu, Scaling: catalog.Scaling{Min: 0, Max: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	rt.SetCatalog(c)
	act := &fakeActivator{onDemand: map[string]bool{"nemo": true}}
	rt.SetActivator(act)
	waitForEndpoint(t, rt)
	return rt, act
}

func TestALoadingModelReportsHowMuchOfItHasArrived(t *testing.T) {
	// The chat page draws its bar from this. Two shards of one replica are one
	// download as far as the person waiting is concerned.
	rt, _ := onDemandRouter(t)
	rt.SetClusterView(replicaSource{replicas: map[string][]adminapi.ReplicaView{
		"nemo": {{State: "starting", AgeSec: 12, Shards: []adminapi.ShardView{
			{Fetch: &adminapi.FetchView{BytesDone: 100, BytesTotal: 400}},
			{Fetch: &adminapi.FetchView{BytesDone: 100, BytesTotal: 400}},
		}}},
	}}, nil)

	st := rt.loadStatus("nemo")
	if st.State != "downloading" || st.Percent == nil || *st.Percent != 25 {
		t.Fatalf("got %+v, expected downloading at 25%%", st)
	}
	if st.Done != 200 || st.Total != 800 {
		t.Errorf("bytes %d of %d, expected 200 of 800", st.Done, st.Total)
	}
}

func TestAModelWithNothingPlacedYetIsWaitingNotStopped(t *testing.T) {
	// Between the wake and the placement there is no replica at all. Calling
	// that "stopped" would end the bar the moment it appeared.
	rt, _ := onDemandRouter(t)
	rt.SetClusterView(replicaSource{}, nil)
	if st := rt.loadStatus("nemo"); st.State != "waiting" {
		t.Errorf("got %+v, expected waiting", st)
	}
}

func TestARunningModelIsReady(t *testing.T) {
	rt, _ := onDemandRouter(t)
	if st := rt.loadStatus("qwen2.5-coder-7b-instruct"); st.State != "ready" {
		t.Errorf("got %+v, expected ready", st)
	}
}

func TestWakingAModelAsksForItAndAnswersAtOnce(t *testing.T) {
	rt, act := onDemandRouter(t)
	rt.SetClusterView(replicaSource{}, nil)
	_, user := chatRouter(t, "http://unused")
	user.ProjectID = ""

	req := httptest.NewRequest("POST", "/chat/wake", strings.NewReader(`{"model":"nemo"}`))
	rec := httptest.NewRecorder()
	rt.handleChatWake(rec, req, user)
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	var st modelStatus
	if err := json.Unmarshal(rec.Body.Bytes(), &st); err != nil {
		t.Fatal(err)
	}
	if st.State != "waiting" {
		t.Errorf("got %+v, expected waiting", st)
	}
	if len(act.asked) != 1 || act.asked[0] != "nemo" {
		t.Errorf("the activator was asked for %v", act.asked)
	}
}

func TestWakingAnUnknownModelIsNotFound(t *testing.T) {
	rt, act := onDemandRouter(t)
	_, user := chatRouter(t, "http://unused")
	user.ProjectID = ""
	req := httptest.NewRequest("POST", "/chat/wake", strings.NewReader(`{"model":"nope"}`))
	rec := httptest.NewRecorder()
	rt.handleChatWake(rec, req, user)
	if rec.Code != 404 {
		t.Errorf("status %d, expected 404", rec.Code)
	}
	if len(act.asked) != 0 {
		t.Errorf("an unknown model was asked for: %v", act.asked)
	}
}