ModelRouter/internal/gateway/sources_api_test.go

package gateway

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

	"llmsproxy/internal/config"
)

// These cover the two write-path defects found while adding a model to a live
// source by hand:
//
//  1. PUT /api/sources/{name} was documented in the API index but not
//     implemented, so it answered 405 while POST (which does the same upsert)
//     worked. A caller following the documented contract could not edit a
//     source at all.
//  2. POST is an upsert that replaces the whole source, so a partial edit that
//     did not carry api_key persisted an empty/placeholder credential. The
//     source kept its name and models and answered 401 on the next call, well
//     after the writing script had exited 0.
//
// The second one is the dangerous class: a write that reports success and
// breaks the source later.

// sourceAPIKey reads the stored key back the way an operator would, i.e. what
// the routes actually use, not what the file happens to contain.
func sourceAPIKey(t *testing.T, g *Gateway, name string) string {
	t.Helper()
	for _, s := range g.core.Sources() {
		if s.Name == name {
			return s.APIKey
		}
	}
	t.Fatalf("source %q not found", name)
	return ""
}

func seededGateway(t *testing.T) (*Gateway, string) {
	t.Helper()
	up := mockUpstream()
	t.Cleanup(up.Close)
	g := newTestGateway(t, config.Source{
		Name:    "seeded",
		BaseURL: up.URL,
		Adapter: "openai",
		APIKey:  "sk-real-credential",
		Models:  []config.Model{{ID: "m-one", Kind: "chat"}},
	})
	return g, up.URL
}

func TestSourcePutUpdatesSource(t *testing.T) {
	g, up := seededGateway(t)

	// A single-field edit: change one model's priority, name the source only in
	// the path. This is the shape the API index documents for PUT.
	rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
		`{"base_url":%q,"adapter":"openai","api_key":"sk-real-credential","models":[{"id":"m-one","priority":42,"kind":"chat"}]}`, up))
	if rr.Code != http.StatusOK {
		t.Fatalf("PUT status=%d body=%s", rr.Code, rr.Body.String())
	}

	var found bool
	for _, s := range g.core.Sources() {
		if s.Name != "seeded" {
			continue
		}
		found = true
		for _, m := range s.Models {
			if m.ID == "m-one" && m.Priority != 42 {
				t.Fatalf("priority not applied: %+v", s.Models)
			}
		}
	}
	if !found {
		t.Fatal("source vanished after PUT")
	}
	// And it must be routable, not merely stored.
	if rr := doReq(t, g, "GET", "/v1/models", ""); rr.Code != http.StatusOK ||
		!strings.Contains(rr.Body.String(), "m-one") {
		t.Fatalf("model not routable after PUT: %d %s", rr.Code, rr.Body.String())
	}
}

func TestSourcePutRejectsNameMismatch(t *testing.T) {
	g, _ := seededGateway(t)
	// Body names a different source than the path. Silently preferring either
	// one would let a client edit a source it did not name.
	rr := doReq(t, g, "PUT", "/api/sources/seeded", `{"name":"other","base_url":"http://x","models":[{"id":"m"}]}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("name mismatch accepted: status=%d body=%s", rr.Code, rr.Body.String())
	}
	if sourceAPIKey(t, g, "seeded") != "sk-real-credential" {
		t.Fatal("rejected PUT still mutated the source")
	}
}

func TestSourcePutRequiresNameInPath(t *testing.T) {
	g, _ := seededGateway(t)
	rr := doReq(t, g, "PUT", "/api/sources", `{"base_url":"http://x","models":[{"id":"m"}]}`)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("PUT without a name in the path: status=%d body=%s", rr.Code, rr.Body.String())
	}
}

func TestSourceUpsertKeepsAPIKey(t *testing.T) {
	for _, tc := range []struct {
		name    string
		sentKey string
	}{
		{"explicit mask", keepExistingAPIKey},
		{"omitted key", ""},
	} {
		t.Run(tc.name, func(t *testing.T) {
			g, up := seededGateway(t)
			rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
				`{"base_url":%q,"adapter":"openai","api_key":%q,"models":[{"id":"m-one"},{"id":"m-two"}]}`,
				up, tc.sentKey))
			if rr.Code != http.StatusOK {
				t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
			}
			if got := sourceAPIKey(t, g, "seeded"); got != "sk-real-credential" {
				t.Fatalf("api_key was clobbered: got %q, want the pre-existing credential", got)
			}
			// The failure this guards is a source that looks configured but
			// cannot authenticate, so assert the route still works.
			if rr := doReq(t, g, "GET", "/v1/models", ""); rr.Code != http.StatusOK {
				t.Fatalf("source unusable after edit: %d %s", rr.Code, rr.Body.String())
			}
		})
	}
}

func TestSourceUpsertStillWritesNewKey(t *testing.T) {
	// The keep-existing path must not swallow a deliberate rotation.
	g, up := seededGateway(t)
	rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
		`{"base_url":%q,"adapter":"openai","api_key":"sk-rotated","models":[{"id":"m-one"}]}`, up))
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	if got := sourceAPIKey(t, g, "seeded"); got != "sk-rotated" {
		t.Fatalf("rotation ignored: got %q", got)
	}
}

func TestSourceUpsertNewSourceMayHaveNoKey(t *testing.T) {
	// Empty is a real value for a credential-less source (local upstreams), so
	// inherit-on-empty must not invent one for a source that does not exist yet.
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t)
	rr := doReq(t, g, "POST", "/api/sources", fmt.Sprintf(
		`{"name":"noauth","base_url":%q,"adapter":"openai","models":[{"id":"local-m"}]}`, up.URL))
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	if got := sourceAPIKey(t, g, "noauth"); got != "" {
		t.Fatalf("unexpected key on a new credential-less source: %q", got)
	}
}

func TestSourceModelIDsIsAdditive(t *testing.T) {
	// "Add one model" is the edit that motivated the defect: it needs no read
	// of the existing list, so it must not require one.
	g, up := seededGateway(t)
	rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
		`{"base_url":%q,"adapter":"openai","api_key":%q,"model_ids":["m-two","m-two"," m-three "]}`, up, keepExistingAPIKey))
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	var ids []string
	for _, s := range g.core.Sources() {
		if s.Name == "seeded" {
			for _, m := range s.Models {
				ids = append(ids, m.ID)
			}
		}
	}
	// m-one kept, m-two not duplicated, blank entries dropped, whitespace trimmed.
	want := map[string]int{"m-one": 1, "m-two": 1, "m-three": 1}
	if len(ids) != len(want) {
		t.Fatalf("models=%v, want exactly %v", ids, want)
	}
	seen := map[string]int{}
	for _, id := range ids {
		seen[id]++
	}
	for id, n := range want {
		if seen[id] != n {
			t.Fatalf("models=%v: %q appears %d times, want %d", ids, id, seen[id], n)
		}
	}
}

func TestSourceWriteRejectsNonAdmin(t *testing.T) {
	g, up := seededGateway(t)
	for _, m := range []string{"POST", "PUT", "DELETE"} {
		req, _ := http.NewRequest(m, "/api/sources/seeded", strings.NewReader(fmt.Sprintf(
			`{"name":"seeded","base_url":%q,"models":[{"id":"m"}]}`, up)))
		req.Header.Set("Authorization", "Bearer sk-not-a-real-key")
		req.Header.Set("Content-Type", "application/json")
		rr := httptest.NewRecorder()
		g.Handler().ServeHTTP(rr, req)
		if rr.Code == http.StatusOK {
			t.Fatalf("%s allowed for an unknown key: %s", m, rr.Body.String())
		}
	}
	if sourceAPIKey(t, g, "seeded") != "sk-real-credential" {
		t.Fatal("unauthenticated write mutated the source")
	}
}

// TestSourceWritePersistsToDisk guards the difference between "the in-memory
// registry accepted it" and "it will survive a restart". The api_key must be
// sealed at rest, never the literal placeholder.
func TestSourceWritePersistsToDisk(t *testing.T) {
	g, up := seededGateway(t)
	rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
		`{"base_url":%q,"adapter":"openai","api_key":%q,"models":[{"id":"m-one"}]}`, up, keepExistingAPIKey))
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	raw, err := os.ReadFile(g.core.Config().Path)
	if err != nil {
		t.Fatalf("read config: %v", err)
	}
	text := string(raw)
	if strings.Contains(text, keepExistingAPIKey) {
		t.Fatal("the keep-mask was persisted as a literal api_key")
	}
	if strings.Contains(text, "sk-real-credential") {
		t.Fatal("api_key stored in plaintext; it must be sealed at rest")
	}
	if !strings.Contains(text, "m-one") {
		t.Fatal("source not persisted")
	}
	// The masked GET view must not leak the credential either.
	rr = doReq(t, g, "GET", "/api/v1/sources/seeded", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("get status=%d", rr.Code)
	}
	var view struct {
		Source struct {
			APIKey string `json:"api_key"`
		} `json:"source"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &view); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if view.Source.APIKey == "sk-real-credential" {
		t.Fatal("GET /api/v1/sources leaked the plaintext api_key")
	}
}