llmgw/internal/router/pricing_test.go

package router

import (
	"crypto/sha256"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

// pricedRouter builds a gateway with authentication, one model under an alias,
// and a price table that charges that model and almost nothing for anything
// else — so a request priced by the wrong name is visible in the figure rather
// than merely plausible.
func pricedRouter(t *testing.T, engineURL string, served ...string) (*Router, string) {
	t.Helper()

	cfg := baseConfig(config.Endpoint{
		ID: "a", URL: engineURL,
		Models: append([]string{"qwen2.5-14b-instruct"}, served...),
	})
	cfg.Models = []catalog.Model{{
		Name: "qwen2.5-14b-instruct", Aliases: []string{"q14"},
	}}
	cfg.Auth.Enabled = true
	cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
	cfg.Auth.Pricing = map[string]config.Price{
		"qwen2.5-14b-instruct": {PromptPer1k: 14000, CompletionPer1k: 42000},
	}
	// Deliberately not zero: a fallback of nothing would make a miss look the
	// same as a model that is free.
	cfg.Auth.DefaultPrice = config.Price{PromptPer1k: 1, CompletionPer1k: 1}
	cfg.Auth.FlushInterval = 20 * time.Millisecond
	cfg.Auth.BatchSize = 1

	rt := newTestRouter(t, cfg)

	org, err := rt.store.CreateOrg(t.Context(), "acme")
	if err != nil {
		t.Fatal(err)
	}
	team, err := rt.store.CreateTeam(t.Context(), org, "platform")
	if err != nil {
		t.Fatal(err)
	}
	project, err := rt.store.CreateProject(t.Context(), team, "gateway")
	if err != nil {
		t.Fatal(err)
	}

	// A real key shape: authentication rejects anything else before it looks.
	secret := store.KeyPrefix + strings.Repeat("a", 43)
	sum := sha256.Sum256([]byte(secret))
	if err := rt.store.CreateKey(t.Context(), store.Key{
		ID: "key-1", Hash: sum[:], Prefix: "aaaa", Name: "test",
		ProjectID: project, CreatedAt: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}
	return rt, secret
}

func TestAnAliasIsPricedAsTheModelItResolvesTo(t *testing.T) {
	// The price table is written against the catalog, and a caller asks for
	// whichever name they were given. Charging by the string they typed meant
	// every alias missed the table and fell to the default — so a budget could
	// be spent at the wrong rate, or not at all, depending on which name went
	// into the request.
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)

	rec := post(t, rt.Handler(),
		map[string]any{"model": "q14", "messages": []any{
			map[string]any{"role": "user", "content": "hello"}}},
		map[string]string{"Authorization": "Bearer " + secret})
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}

	u := waitForUsage(t, rt)
	want := charge(u.Prompt, 14000) + charge(u.Completion, 42000)
	fallback := charge(u.Prompt, 1) + charge(u.Completion, 1)
	if u.CostMicros != want {
		t.Errorf("charged %d micros for %d+%d tokens, expected %d; the "+
			"fallback rate would have charged %d",
			u.CostMicros, u.Prompt, u.Completion, want, fallback)
	}
}

// charge is Price.Cost for one half of a request, rounding up the way the
// accountant does — spelled out here so the expected figure is arithmetic the
// test states rather than the code under test restated.
func charge(tokens, per1k int64) int64 {
	if tokens <= 0 {
		return 0
	}
	return (tokens*per1k + 999) / 1000
}

func TestAModelWithNoPriceFallsBackRatherThanCostingNothing(t *testing.T) {
	e := newFakeEngine(t, "unpriced")
	rt, secret := pricedRouter(t, e.srv.URL, "unpriced")

	// Not in the catalog and not in the price table. It must still cost the
	// fallback, because a model nobody priced is the one most likely to be new.
	rec := post(t, rt.Handler(),
		map[string]any{"model": "unpriced", "messages": []any{
			map[string]any{"role": "user", "content": "hello"}}},
		map[string]string{"Authorization": "Bearer " + secret})
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}

	u := waitForUsage(t, rt)
	want := charge(u.Prompt, 1) + charge(u.Completion, 1)
	if u.CostMicros != want {
		t.Errorf("charged %d micros for %d+%d tokens, expected the fallback's %d",
			u.CostMicros, u.Prompt, u.Completion, want)
	}
	if u.CostMicros == 0 {
		t.Error("an unpriced model cost nothing; the fallback did not apply")
	}
}

func waitForUsage(t *testing.T, rt *Router) store.KeyUsage {
	t.Helper()
	since := time.Now().Add(-time.Hour)
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		rows, err := rt.store.UsageSince(t.Context(), since)
		if err == nil && len(rows) > 0 {
			return rows[0]
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("no usage was recorded")
	return store.KeyUsage{}
}