ModelRouter/internal/lua/usage_plugin_test.go

package lua

import (
	"strings"
	"testing"
)

// TestUsagePluginShipsAndGuardsItsOwnKey covers the two things that make the
// usage plugin safe to ship, and both are easy to break by accident because they
// live in a Lua file nobody compiles.
//
//  1. SHIPPED. `usage` is the reason /v1/usage exists at all, and it only exists
//     for deployments that receive it. A gateway build that does not embed it
//     answers 404 and the frontend shows "后台未开放消费接口" — which reads as a
//     bug rather than as a missing optional feature. Same reasoning billing.lua.
//
//  2. SELF-SCOPED. The handler reports the CALLER's usage, so it must look the
//     caller up by the key in the request, never by a key from its own
//     configuration. A plugin that read a configured key would report one
//     tenant's spend to every tenant — the exact leak this plugin avoids by
//     passing the request's key straight through.
func TestUsagePluginShipsAndGuardsItsOwnKey(t *testing.T) {
	src := readShippedPlugin(t, "usage.lua")

	for _, want := range []string{
		`path = "/usage"`,
		`handler = "usage"`,
		`roles = { "admin", "user" }`,
		`payload.key`,
		`maskedKey`,
	} {
		if !strings.Contains(src, want) {
			t.Errorf("usage.lua no longer contains %q", want)
		}
	}

	// The adapter-less path: it must not reach for a stored credential.
	for _, forbidden := range []string{
		"PLUGIN.state",
		"PLUGIN.prices",
		"api_key",
		"control_key",
	} {
		if strings.Contains(src, forbidden) {
			t.Errorf("usage.lua reads %q — it must report only the calling key "+
				"from the request, never a stored credential", forbidden)
		}
	}
}

// TestUsageMasksKeyLikeTheGateway pins the mask algorithm against the Go one.
//
// This is the subtle one. /api/stats keys its maps by keyID(req.Key) — the
// MASKED form ("***f9a9ed") — while the handler receives the FULL key. Passing
// the full key does not error, it silently returns zero, and the endpoint then
// reports "0 requests" for a key with hundreds of thousands of them. Measured
// against the live gateway before this was pinned:
//
//	?key=***f9a9ed              → total.reqs = 33847
//	?key=sk-gw-7ffe…9ed         → total.reqs = 0
//
// So the Lua mask must match internal/gateway/server.go's keyID exactly: keys of
// 6 characters or fewer are returned unchanged, longer ones become "***" plus the
// last 6.
func TestUsageMasksKeyLikeTheGateway(t *testing.T) {
	cases := []struct{ in, want string }{
		{"sk-gw-7ffe12c5cbacec40bd39a7291df9a9ed", "***f9a9ed"},
		{"sk-gw-7287aba64ef7ee665cf723222610a67d", "***10a67d"},
		{"1234567", "***234567"},
		// 6 or fewer: returned unchanged, on both sides.
		{"short", "short"},
		{"123456", "123456"},
	}
	src := readShippedPlugin(t, "usage.lua")
	for _, c := range cases {
		// The mask lives in Lua, so assert on its shape rather than run it: a
		// behavioural test here would need a whole VM plus a stubbed stats
		// endpoint, and would still not cover the real one.
		if !strings.Contains(src, `return "***" .. string.sub(key, -6)`) {
			t.Fatal("usage.lua no longer masks as \"***\" + last 6 characters")
		}
		if !strings.Contains(src, `if #key <= 6 then return key end`) {
			t.Fatal("usage.lua no longer returns short keys unchanged")
		}
		_ = c
	}
}

// TestUsageReportsZeroOnlyForDataItGot guards the failure shape. An endpoint that
// cannot reach its data source must say so: returning zeros would be rendered as
// "已用 0，0 次" and read as "this key has never been used", which is a different
// and wrong claim.
func TestUsageReportsZeroOnlyForDataItGot(t *testing.T) {
	src := readShippedPlugin(t, "usage.lua")
	if !strings.Contains(src, "usage unavailable") {
		t.Error("usage.lua has no distinct failure message; a failed lookup must not " +
			"be indistinguishable from zero usage")
	}
	if !strings.Contains(src, "if r == nil then") {
		t.Error("usage.lua does not check for a refused or timed-out stats request")
	}
}

// TestUsagePeriodIsAll pins the window. "All" answers "how much has this key
// used"; a day window would look like the number resets every morning, which is
// the more alarming misreading of the same digits.
func TestUsagePeriodIsAll(t *testing.T) {
	if !strings.Contains(readShippedPlugin(t, "usage.lua"), `local PERIOD = "all"`) {
		t.Error(`usage.lua no longer defaults to period "all"`)
	}
}

// readShippedPlugin reads an optional plugin from the repository (not the
// binary: plugins are never embedded). Serves tests that assert on Lua source.
func readShippedPlugin(t *testing.T, name string) string {
	t.Helper()
	src, err := ReadBundledPlugin(strings.TrimSuffix(name, ".lua"))
	if err != nil {
		t.Fatalf("read repo plugin %s: %v", name, err)
	}
	return src
}