116 lines
3.3 KiB
Go
116 lines
3.3 KiB
Go
package fs
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"ollie/cmd/olliesrv/internal/backend"
|
|
)
|
|
|
|
// stubCtlBackend satisfies backend.Backend for pricing-block tests.
|
|
type stubCtlBackend struct{}
|
|
|
|
func (stubCtlBackend) ChatStream(context.Context, []backend.Message, []backend.Tool, backend.GenerationParams) (<-chan backend.StreamEvent, error) {
|
|
return nil, fmt.Errorf("not implemented")
|
|
}
|
|
func (stubCtlBackend) Name() string { return "stub" }
|
|
func (stubCtlBackend) DefaultModel() string { return "stub-model" }
|
|
func (stubCtlBackend) Model() string { return "stub-model" }
|
|
func (stubCtlBackend) SetModel(string) {}
|
|
func (stubCtlBackend) ContextLength(context.Context) int { return 0 }
|
|
func (stubCtlBackend) Models(context.Context) []string { return nil }
|
|
|
|
// stubCtlLister adds API-provided pricing for one exact model ID.
|
|
type stubCtlLister struct{ stubCtlBackend }
|
|
|
|
func (stubCtlLister) ModelsInfo(context.Context) []backend.ModelInfo {
|
|
return []backend.ModelInfo{
|
|
{ID: "exact-model", Pricing: &backend.ModelPricing{
|
|
Input: 2.5 / 1_000_000, Output: 10.0 / 1_000_000, CacheRead: 1.25 / 1_000_000,
|
|
}},
|
|
}
|
|
}
|
|
|
|
func TestModelPricingBlock(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
be backend.Backend
|
|
id string
|
|
want string
|
|
}{
|
|
{
|
|
name: "api_pricing_exact_match",
|
|
be: stubCtlLister{},
|
|
id: "exact-model",
|
|
want: "exact-model\n" +
|
|
"In/Mtok: $2.50\n" +
|
|
"Out/Mtok: $10.00\n" +
|
|
"CacheRead/Mtok: $1.25\n" +
|
|
"CacheWrite/Mtok: -\n",
|
|
},
|
|
{
|
|
name: "api_pricing_miss_falls_back_to_static",
|
|
be: stubCtlLister{},
|
|
id: "claude-sonnet-4-5",
|
|
want: "claude-sonnet-4-5\n" +
|
|
"In/Mtok: $3.00(e)\n" +
|
|
"Out/Mtok: $15.00(e)\n" +
|
|
"CacheRead/Mtok: $0.3000(e)\n" +
|
|
"CacheWrite/Mtok: $3.75(e)\n",
|
|
},
|
|
{
|
|
name: "static_pricing_no_lister",
|
|
be: stubCtlBackend{},
|
|
id: "gpt-4o",
|
|
want: "gpt-4o\n" +
|
|
"In/Mtok: $2.50(e)\n" +
|
|
"Out/Mtok: $10.00(e)\n" +
|
|
"CacheRead/Mtok: $1.25(e)\n" +
|
|
"CacheWrite/Mtok: -\n",
|
|
},
|
|
{
|
|
name: "no_pricing_known",
|
|
be: stubCtlBackend{},
|
|
id: "totally-unknown-model",
|
|
want: "totally-unknown-model\n",
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := modelPricingBlock(tt.be, tt.id)
|
|
if got != tt.want {
|
|
t.Errorf("modelPricingBlock(%q) =\n%q\nwant\n%q", tt.id, got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFormatPricingWithEstimate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
perToken float64
|
|
estimated bool
|
|
want string
|
|
}{
|
|
{"zero", 0, false, "-"},
|
|
{"zero_estimated", 0, true, "-"},
|
|
{"sonnet_input", 3.0 / 1_000_000, false, "$3.00"},
|
|
{"sonnet_input_estimated", 3.0 / 1_000_000, true, "$3.00(e)"},
|
|
{"sonnet_output", 15.0 / 1_000_000, false, "$15.00"},
|
|
{"sonnet_cache_read", 0.30 / 1_000_000, false, "$0.3000"},
|
|
{"sonnet_cache_read_estimated", 0.30 / 1_000_000, true, "$0.3000(e)"},
|
|
{"haiku_input", 1.0 / 1_000_000, false, "$1.00"},
|
|
{"small_value", 0.01 / 1_000_000, false, "$0.0100"},
|
|
{"small_value_estimated", 0.01 / 1_000_000, true, "$0.0100(e)"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := formatPricingWithEstimate(tt.perToken, tt.estimated)
|
|
if got != tt.want {
|
|
t.Errorf("formatPricingWithEstimate(%v, %v) = %q, want %q", tt.perToken, tt.estimated, got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|