Vectrum
Sign In
72 lines · 2.1 KB · Go
package llm

import (
	"strings"
	"testing"
)

// The model picker offers whatever /v1/models returns, which is dated ids and
// context-window variants. Those are the same model at the same price, so
// pricing has to look past the decoration rather than report "unknown".
func TestCostPricesDecoratedModelIDs(t *testing.T) {
	u := Usage{InputTokens: 1_000_000}
	base, ok := u.Cost("claude-sonnet-5")
	if !ok {
		t.Fatal("the bare id is not priced")
	}
	for _, id := range []string{
		"claude-sonnet-5-20250930",
		"claude-sonnet-5[1m]",
		"claude-sonnet-5-latest",
	} {
		got, ok := u.Cost(id)
		if !ok {
			t.Errorf("Cost(%q) reported no price", id)
			continue
		}
		if got != base {
			t.Errorf("Cost(%q) = %v, want the same %v as the bare id", id, got, base)
		}
	}

	// A model nobody has priced still reports nothing rather than guessing.
	if _, ok := u.Cost("some-gateway/mystery-model"); ok {
		t.Error("an unknown model was given a price")
	}
	if _, ok := u.Cost("claude-something-9-20990101"); ok {
		t.Error("an unknown Claude model was given a price")
	}
}

// The wizard's fallback list has to stand on its own: everything in it priced,
// the default present, and the same order every time.
func TestPricedModelsIsAUsableFallback(t *testing.T) {
	got := PricedModels()
	if len(got) != len(prices) {
		t.Fatalf("PricedModels returned %d of %d priced models", len(got), len(prices))
	}

	var hasDefault bool
	for _, id := range got {
		if id == DefaultModel {
			hasDefault = true
		}
		if _, ok := PriceOf(id); !ok {
			t.Errorf("%s is offered but cannot be priced", id)
		}
	}
	if !hasDefault {
		t.Errorf("the fallback omits the default model %s", DefaultModel)
	}
	if first := got[0]; !strings.HasPrefix(first, "claude-opus") {
		t.Errorf("the list opens on %s, want the flagship family first", first)
	}

	// Map iteration is random, so an unsorted implementation passes once in a
	// while and fails in front of a user.
	for i := 0; i < 20; i++ {
		if again := PricedModels(); strings.Join(again, ",") != strings.Join(got, ",") {
			t.Fatalf("the order is not stable:\n%v\n%v", got, again)
		}
	}
}