CodeRabbit Generated Unit Tests: Generate unit tests for PR changes

2026-06-16 12:47:43 +00:00
parent 6cf1317341
commit 090bb4623c
4 changed files with 521 additions and 2 deletions
@@ -1,6 +1,7 @@
 package scheduler

 import (
+	"context"
 	"errors"
 	"io"
 	"net/http"
@@ -54,8 +55,9 @@ type stopRec struct {
 // fakeEffects is an in-memory scheduler.Effects. Tests program process states
 // and GrantServe outcomes, then assert on the recorded calls.
 type fakeEffects struct {
-	states      map[string]process.ProcessState // model -> state; missing => not handled
-	serveResult map[string]bool                 // GrantServe return per model (default true)
+	states       map[string]process.ProcessState // model -> state; missing => not handled
+	serveResult  map[string]bool                 // GrantServe return per model (default true)
+	lastServeReq HandlerReq

 	starts []startRec
 	grants []grantRec
@@ -98,6 +100,7 @@ func (f *fakeEffects) GrantServe(req HandlerReq, modelID string) bool {
 	if v, set := f.serveResult[modelID]; set {
 		ok = v
 	}
+	f.lastServeReq = req
 	f.grants = append(f.grants, grantRec{model: modelID, serve: ok})
 	return ok
 }
@@ -169,6 +172,99 @@ func TestFIFO_FastPath(t *testing.T) {
 	}
 }

+func TestFIFO_GrantSetsPriorityMetadata(t *testing.T) {
+	eff := newFakeEffects()
+	eff.states["a"] = process.StateReady
+	cfg := config.FifoConfig{Priority: map[string]int{"a": 7}}
+	s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff)
+
+	ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a", Metadata: make(map[string]string)})
+	s.OnRequest(HandlerReq{Model: "a", Ctx: ctx})
+
+	if got := eff.served("a"); got != 1 {
+		t.Fatalf("served(a)=%d want 1", got)
+	}
+	data, ok := shared.ReadContext(eff.lastServeReq.Ctx)
+	if !ok {
+		t.Fatal("context data missing from granted request")
+	}
+	if data.Metadata["fifo_priority"] != "7" {
+		t.Errorf("fifo_priority = %q, want 7", data.Metadata["fifo_priority"])
+	}
+}
+
+func TestFIFO_GrantSetsPriorityMetadata_DefaultZero(t *testing.T) {
+	// A model that is not listed in the Priority map should get fifo_priority="0".
+	eff := newFakeEffects()
+	eff.states["unlisted"] = process.StateReady
+	cfg := config.FifoConfig{Priority: map[string]int{"other": 5}} // "unlisted" absent
+	s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff)
+
+	ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "unlisted", Metadata: make(map[string]string)})
+	s.OnRequest(HandlerReq{Model: "unlisted", Ctx: ctx})
+
+	if got := eff.served("unlisted"); got != 1 {
+		t.Fatalf("served(unlisted)=%d want 1", got)
+	}
+	data, ok := shared.ReadContext(eff.lastServeReq.Ctx)
+	if !ok {
+		t.Fatal("context data missing from granted request")
+	}
+	if data.Metadata["fifo_priority"] != "0" {
+		t.Errorf("fifo_priority = %q, want %q", data.Metadata["fifo_priority"], "0")
+	}
+}
+
+func TestFIFO_GrantSetsPriorityMetadata_NoMetadataMap(t *testing.T) {
+	// When the request context has no Metadata map, grantHandler must not crash.
+	// It should log a debug message and still grant the request.
+	eff := newFakeEffects()
+	eff.states["a"] = process.StateReady
+	cfg := config.FifoConfig{Priority: map[string]int{"a": 3}}
+	s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff)
+
+	// No Metadata map in the context data — SetReqData will return an error.
+	ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a"})
+	s.OnRequest(HandlerReq{Model: "a", Ctx: ctx})
+
+	// The grant must still succeed despite the missing metadata map.
+	if got := eff.served("a"); got != 1 {
+		t.Fatalf("served(a)=%d want 1 (metadata error must not prevent grant)", got)
+	}
+}
+
+func TestFIFO_GrantSetsPriorityMetadata_AfterSwapCompletion(t *testing.T) {
+	// Priority metadata must be set for waiters granted via OnSwapDone, not just
+	// requests that hit the fast path.
+	eff := newFakeEffects()
+	eff.states["a"] = process.StateStopped // forces a swap
+	cfg := config.FifoConfig{Priority: map[string]int{"a": 9}}
+	s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff)
+
+	ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a", Metadata: make(map[string]string)})
+	s.OnRequest(HandlerReq{Model: "a", Ctx: ctx})
+
+	// Swap is in flight; no grant yet.
+	if got := eff.served("a"); got != 0 {
+		t.Fatalf("served(a)=%d want 0 before swap done", got)
+	}
+
+	// Complete the swap.
+	eff.states["a"] = process.StateReady
+	s.OnSwapDone(SwapDone{ModelID: "a"})
+
+	if got := eff.served("a"); got != 1 {
+		t.Fatalf("served(a)=%d want 1 after swap done", got)
+	}
+	data, ok := shared.ReadContext(eff.lastServeReq.Ctx)
+	if !ok {
+		t.Fatal("context data missing from granted request after swap")
+	}
+	if data.Metadata["fifo_priority"] != "9" {
+		t.Errorf("fifo_priority = %q, want %q", data.Metadata["fifo_priority"], "9")
+	}
+}
+
 func TestFIFO_ModelNotFound(t *testing.T) {
 	eff := newFakeEffects() // no states => model unknown
 	s := newFIFO(&stubPlanner{}, eff)
@@ -1,9 +1,13 @@
 package server

 import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"

+	"github.com/mostlygeek/llama-swap/internal/shared"
 	"github.com/tidwall/gjson"
 )

@@ -56,6 +60,109 @@ func TestServer_ProcessStreamingResponse_NoData(t *testing.T) {
 	}
 }

+func TestMetricsMonitor_RecordMetadata(t *testing.T) {
+	mm := newMetricsMonitor(nil, 10, 0)
+	r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"usage":{}}`))
+	r = r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{
+		ModelID:  "m",
+		Metadata: map[string]string{"client": "web", "trace": "abc"},
+	}))
+
+	w := httptest.NewRecorder()
+	copier := newBodyCopier(w)
+	copier.WriteHeader(http.StatusOK)
+	copier.Write([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2}}`))
+
+	mm.record("m", r, copier, 0, nil, nil)
+
+	entries := mm.getMetrics()
+	if len(entries) != 1 {
+		t.Fatalf("want 1 entry, got %d", len(entries))
+	}
+	if entries[0].Metadata["client"] != "web" {
+		t.Errorf("client = %q, want web", entries[0].Metadata["client"])
+	}
+	if entries[0].Metadata["trace"] != "abc" {
+		t.Errorf("trace = %q, want abc", entries[0].Metadata["trace"])
+	}
+}
+
+func TestMetricsMonitor_RecordMetadata_EmptyMap(t *testing.T) {
+	// An empty Metadata map in context must NOT set tm.Metadata (omitempty semantics).
+	mm := newMetricsMonitor(nil, 10, 0)
+	r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{}`))
+	r = r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{
+		ModelID:  "m",
+		Metadata: map[string]string{}, // empty, not nil
+	}))
+
+	w := httptest.NewRecorder()
+	copier := newBodyCopier(w)
+	copier.WriteHeader(http.StatusOK)
+	copier.Write([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2}}`))
+
+	mm.record("m", r, copier, 0, nil, nil)
+
+	entries := mm.getMetrics()
+	if len(entries) != 1 {
+		t.Fatalf("want 1 entry, got %d", len(entries))
+	}
+	if entries[0].Metadata != nil {
+		t.Errorf("Metadata should be nil for empty context metadata, got %v", entries[0].Metadata)
+	}
+}
+
+func TestMetricsMonitor_RecordMetadata_NoContextData(t *testing.T) {
+	// A request with no ReqContextData in context should produce nil Metadata.
+	mm := newMetricsMonitor(nil, 10, 0)
+	r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{}`))
+	// No shared.SetContext call — no ReqContextData in context.
+
+	w := httptest.NewRecorder()
+	copier := newBodyCopier(w)
+	copier.WriteHeader(http.StatusOK)
+	copier.Write([]byte(`{"usage":{"prompt_tokens":3,"completion_tokens":4}}`))
+
+	mm.record("m", r, copier, 0, nil, nil)
+
+	entries := mm.getMetrics()
+	if len(entries) != 1 {
+		t.Fatalf("want 1 entry, got %d", len(entries))
+	}
+	if entries[0].Metadata != nil {
+		t.Errorf("Metadata should be nil when no context data, got %v", entries[0].Metadata)
+	}
+}
+
+func TestMetricsMonitor_RecordMetadata_DeepCopy(t *testing.T) {
+	// Mutating the original context metadata after record() must not affect the stored entry.
+	mm := newMetricsMonitor(nil, 10, 0)
+	original := map[string]string{"key": "before"}
+	r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{}`))
+	r = r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{
+		ModelID:  "m",
+		Metadata: original,
+	}))
+
+	w := httptest.NewRecorder()
+	copier := newBodyCopier(w)
+	copier.WriteHeader(http.StatusOK)
+	copier.Write([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2}}`))
+
+	mm.record("m", r, copier, 0, nil, nil)
+
+	// Mutate the original map after record.
+	original["key"] = "after"
+
+	entries := mm.getMetrics()
+	if len(entries) != 1 {
+		t.Fatalf("want 1 entry, got %d", len(entries))
+	}
+	if entries[0].Metadata["key"] != "before" {
+		t.Errorf("Metadata[key] = %q, want %q (deep copy expected)", entries[0].Metadata["key"], "before")
+	}
+}
+
 func TestServer_ParseMetrics_Infill(t *testing.T) {
 	// /infill responses are arrays; timings live in the last element.
 	body := `[{"content":"a"},{"content":"b","timings":{"prompt_n":5,"predicted_n":9,"prompt_ms":10,"predicted_ms":20}}]`
@@ -387,6 +387,105 @@ func TestExtractContext_ApiKey(t *testing.T) {
 	}
 }

+func TestSetReqData(t *testing.T) {
+	ctx := SetContext(context.Background(), ReqContextData{Model: "llama3", ModelID: "llama3", Metadata: make(map[string]string)})
+
+	if err := SetReqData(ctx, "client", "web"); err != nil {
+		t.Fatalf("SetReqData: %v", err)
+	}
+	if err := SetReqData(ctx, "trace", "abc123"); err != nil {
+		t.Fatalf("SetReqData: %v", err)
+	}
+
+	data, ok := ReadContext(ctx)
+	if !ok {
+		t.Fatal("context data missing")
+	}
+	if data.Metadata["client"] != "web" {
+		t.Errorf("client = %q, want %q", data.Metadata["client"], "web")
+	}
+	if data.Metadata["trace"] != "abc123" {
+		t.Errorf("trace = %q, want %q", data.Metadata["trace"], "abc123")
+	}
+}
+
+func TestSetReqData_Errors(t *testing.T) {
+	if err := SetReqData(context.Background(), "k", "v"); err == nil {
+		t.Error("expected error when no request context data exists")
+	}
+	ctx := SetContext(context.Background(), ReqContextData{Model: "llama3", ModelID: "llama3"})
+	if err := SetReqData(ctx, "k", "v"); err == nil {
+		t.Error("expected error when metadata map is missing")
+	}
+}
+
+func TestSetReqData_NilContext(t *testing.T) {
+	// nil context must return an error without panicking.
+	err := SetReqData(nil, "k", "v")
+	if err == nil {
+		t.Error("expected error for nil context, got nil")
+	}
+}
+
+func TestSetReqData_OverwritesExistingKey(t *testing.T) {
+	ctx := SetContext(context.Background(), ReqContextData{
+		Model:    "m",
+		Metadata: map[string]string{"key": "old"},
+	})
+	if err := SetReqData(ctx, "key", "new"); err != nil {
+		t.Fatalf("SetReqData: %v", err)
+	}
+	data, _ := ReadContext(ctx)
+	if data.Metadata["key"] != "new" {
+		t.Errorf("key = %q, want %q", data.Metadata["key"], "new")
+	}
+}
+
+func TestExtractContext_MetadataInitialized_GET(t *testing.T) {
+	r, _ := http.NewRequest(http.MethodGet, "/?model=llama3", nil)
+	got, err := extractContext(r)
+	if err != nil {
+		t.Fatalf("extractContext: %v", err)
+	}
+	if got.Metadata == nil {
+		t.Error("Metadata should be initialized (not nil) for GET requests")
+	}
+}
+
+func TestExtractContext_MetadataInitialized_JSON(t *testing.T) {
+	r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"llama3"}`))
+	r.Header.Set("Content-Type", "application/json")
+	got, err := extractContext(r)
+	if err != nil {
+		t.Fatalf("extractContext: %v", err)
+	}
+	if got.Metadata == nil {
+		t.Error("Metadata should be initialized (not nil) for JSON POST requests")
+	}
+}
+
+func TestExtractContext_MetadataInitialized_Form(t *testing.T) {
+	r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", strings.NewReader("model=whisper-1"))
+	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	got, err := extractContext(r)
+	if err != nil {
+		t.Fatalf("extractContext: %v", err)
+	}
+	if got.Metadata == nil {
+		t.Error("Metadata should be initialized (not nil) for form POST requests")
+	}
+}
+
+func TestExtractContext_MetadataIsWritable(t *testing.T) {
+	// Verify the initialized map is writable — i.e. SetReqData can use it.
+	r, _ := http.NewRequest(http.MethodGet, "/?model=llama3", nil)
+	got, _ := extractContext(r)
+	ctx := SetContext(context.Background(), got)
+	if err := SetReqData(ctx, "x", "y"); err != nil {
+		t.Fatalf("SetReqData on extractContext Metadata: %v", err)
+	}
+}
+
 func TestServer_ExtractAPIKey(t *testing.T) {
 	basicHeader := func(user, pass string) string {
 		return "Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+pass))