internal/budget/store_mongodb_test.go
Rama main · réplica de lectura del repositorio original.
package budget
import (
"context"
"errors"
"testing"
"time"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
"github.com/enterpilot/gomodel/internal/storage/mongotest"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestIsMongoTransactionCapabilityError(t *testing.T) {
tests := []struct {
name string
err error
want bool
}{
{
name: "standalone transaction message",
err: errors.New("Transaction numbers are only allowed on a replica set member or mongos"),
want: true,
},
{
name: "illegal operation command code",
err: mongo.CommandError{
Code: 20,
Message: "transaction is not supported by this deployment",
Labels: []string{"TransientTransactionError"},
},
want: true,
},
{
name: "labeled unsupported transaction message",
err: mongo.CommandError{
Message: "transaction is not supported by this deployment",
Labels: []string{"TransientTransactionError"},
},
want: true,
},
{
name: "ordinary transient transaction error",
err: mongo.CommandError{
Message: "temporary write conflict",
Labels: []string{"TransientTransactionError"},
},
want: false,
},
{
name: "ordinary error",
err: errors.New("network timeout"),
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := isMongoTransactionCapabilityError(tt.err)
require.Equal(t, tt.want, got)
})
}
}
func TestUsagePathRegexBoundaries(t *testing.T) {
// The root pattern must also catch the rows a blank user path normalizes
// to; nested paths must not leak into a sibling that shares their prefix.
require.Equal(t, "^/", usagePathRegex("/"))
require.Equal(t, `^/team(?:/|$)`, usagePathRegex("/team"))
}
func TestMongoSubjectMatchUsesLabelMembership(t *testing.T) {
got, err := mongoSubjectMatch(SpendWindow{Scope: ScopeLabel, Subject: "iOS"})
require.NoError(t, err)
want := mongoConsumerSpendMatch(bson.D{{Key: "$in", Value: bson.A{
bson.D{{Key: "$literal", Value: "iOS"}},
bson.D{{Key: "$ifNull", Value: bson.A{"$labels", bson.A{}}}},
}}})
require.Equal(t, want, got)
}
// A subject that starts with "$" would otherwise be read as a field path
// inside the aggregation expression instead of compared as text.
func TestMongoSubjectMatchTreatsDollarSubjectsAsLiterals(t *testing.T) {
tests := []struct {
name string
window SpendWindow
want bson.D
}{
{
name: "label",
window: SpendWindow{Scope: ScopeLabel, Subject: "$labels"},
want: mongoConsumerSpendMatch(bson.D{{Key: "$in", Value: bson.A{
bson.D{{Key: "$literal", Value: "$labels"}},
bson.D{{Key: "$ifNull", Value: bson.A{"$labels", bson.A{}}}},
}}}),
},
{
name: "provider",
window: SpendWindow{Scope: ScopeProvider, Subject: "$provider"},
want: bson.D{{Key: "$eq", Value: bson.A{
bson.D{{Key: "$toLower", Value: mongoTrimmedOr("$provider_name", "$provider")}},
bson.D{{Key: "$literal", Value: "$provider"}},
}}},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := mongoSubjectMatch(tt.window)
require.NoError(t, err)
assert.Equal(t, tt.want, got)
})
}
}
func TestMongoSubjectMatchRejectsBlankLabel(t *testing.T) {
_, err := mongoSubjectMatch(SpendWindow{Scope: ScopeLabel, Subject: " "})
require.Error(t, err)
}
func TestMongoSubjectMatchNormalizesMissingUserPathToRoot(t *testing.T) {
got, err := mongoSubjectMatch(SpendWindow{Scope: ScopeUserPath, Subject: "/team"})
require.NoError(t, err)
// A consumer predicate is the path match and the billable guard.
conjunction, ok := bsonField(got, "$and").(bson.A)
require.True(t, ok, "mongoSubjectMatch(user_path) = %#v, want an $and expression", got)
require.Len(t, conjunction, 2)
pathMatch, ok := conjunction[0].(bson.D)
require.True(t, ok)
expression, ok := bsonField(pathMatch, "$regexMatch").(bson.D)
require.True(t, ok, "mongoSubjectMatch(user_path) = %#v, want a $regexMatch expression", got)
regex := bsonField(expression, "regex")
require.Equal(t, usagePathRegex("/team"), regex)
// The input must fold missing and blank paths to "/" so a root budget still
// sees rows recorded without a user path.
_, ok = bsonField(expression, "input").(bson.D)
require.True(t, ok, "input = %#v, want the normalizing $let expression", bsonField(expression, "input"))
}
func TestMongoDBStoreRoundTripsPerChild(t *testing.T) {
mongotest.Run(t, func(t *testing.T, db *mongo.Database) {
ctx := context.Background()
store, err := NewMongoDBStore(ctx, db)
require.NoError(t, err)
budgets := []Budget{
{Scope: ScopeUserPath, Subject: "/shared", PeriodSeconds: PeriodDailySeconds, Amount: 10, Source: SourceManual},
{Scope: ScopeUserPath, Subject: "/customers", PerChild: true, PeriodSeconds: PeriodDailySeconds, Amount: 20, Source: SourceManual},
}
err = store.UpsertBudgets(ctx, budgets)
require.NoError(t, err)
got, err := store.ListBudgets(ctx)
require.NoError(t, err)
bySubject := make(map[string]Budget, len(got))
for _, budget := range got {
bySubject[budget.Subject] = budget
}
for _, want := range budgets {
persisted, ok := bySubject[want.Subject]
assert.True(t, ok)
assert.Equal(t, want.PerChild, persisted.PerChild, "budget %q = %+v, want per_child=%v", want.Subject, persisted, want.PerChild)
}
})
}
func bsonField(document bson.D, key string) any {
for _, element := range document {
if element.Key == key {
return element.Value
}
}
return nil
}
func TestBsonNumberReadsEveryNumericShape(t *testing.T) {
tests := []struct {
name string
value any
want float64
ok bool
}{
{name: "double", value: 1.5, want: 1.5, ok: true},
{name: "int64", value: int64(3), want: 3, ok: true},
{name: "int32", value: int32(2), want: 2, ok: true},
{name: "missing", value: nil},
{name: "unexpected type", value: "7"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, ok := bsonNumber(tt.value)
require.Equal(t, tt.want, got)
require.Equal(t, tt.ok, ok)
})
}
}
// A MongoDB database written before budget scopes existed carries a unique
// index on (user_path, period_seconds). Unsetting user_path collapses every
// migrated document of the same period onto one index key, so the legacy index
// has to go before the rewrite rather than after it.
func TestMongoDBStoreMigratesPreScopeDocuments(t *testing.T) {
mongotest.Run(t, func(t *testing.T, db *mongo.Database) {
ctx := context.Background()
budgets := db.Collection("budgets")
_, err := budgets.Indexes().CreateOne(ctx, mongo.IndexModel{
Keys: bson.D{{Key: "user_path", Value: 1}, {Key: "period_seconds", Value: 1}},
Options: options.Index().SetUnique(true),
})
require.NoError(t, err)
now := time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)
// Two paths sharing one period: the case that collides.
_, err = budgets.InsertMany(ctx, []any{
bson.D{{Key: "user_path", Value: "/team/alpha"}, {Key: "period_seconds", Value: PeriodDailySeconds},
{Key: "amount", Value: 10.0}, {Key: "source", Value: SourceManual},
{Key: "created_at", Value: now}, {Key: "updated_at", Value: now}},
bson.D{{Key: "user_path", Value: "/team/beta"}, {Key: "period_seconds", Value: PeriodDailySeconds},
{Key: "amount", Value: 20.0}, {Key: "source", Value: SourceManual},
{Key: "created_at", Value: now}, {Key: "updated_at", Value: now}},
})
require.NoError(t, err)
store, err := NewMongoDBStore(ctx, db)
require.NoError(t, err)
got, err := store.ListBudgets(ctx)
require.NoError(t, err)
require.Len(t, got, 2)
bySubject := map[string]Budget{}
for _, budget := range got {
require.Equal(t, ScopeUserPath, budget.Scope, "migrated budget %+v, want scope user_path", budget)
bySubject[budget.Subject] = budget
}
require.Equal(t, float64(10), bySubject["/team/alpha"].Amount)
require.Equal(t, float64(20), bySubject["/team/beta"].Amount, "migrated budgets = %+v, want both amounts preserved", got)
// The scoped unique index must now allow a label budget spelled like an
// existing user path.
err = store.UpsertBudgets(ctx, []Budget{
{Scope: ScopeLabel, Subject: "/team/alpha", PeriodSeconds: PeriodDailySeconds, Amount: 1},
})
require.NoError(t, err)
})
}
func TestMongoDBStoreSumSpendMatchesProvider(t *testing.T) {
mongotest.Run(t, func(t *testing.T, db *mongo.Database) {
ctx := context.Background()
store, err := NewMongoDBStore(ctx, db)
require.NoError(t, err)
now := time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)
_, err = db.Collection("usage").InsertMany(ctx, []any{
bson.D{{Key: "timestamp", Value: now}, {Key: "provider", Value: "openai"}, {Key: "provider_name", Value: "OpenAI-Prod"}, {Key: "total_cost", Value: 2.0}},
bson.D{{Key: "timestamp", Value: now}, {Key: "provider", Value: "openai"}, {Key: "provider_name", Value: "openai-free"}, {Key: "total_cost", Value: 4.0}},
// Written before provider_name existed: only the provider type is set.
bson.D{{Key: "timestamp", Value: now}, {Key: "provider", Value: "openai"}, {Key: "total_cost", Value: 3.0}},
bson.D{{Key: "timestamp", Value: now}, {Key: "provider", Value: "openai"}, {Key: "provider_name", Value: "openai-prod"}, {Key: "cache_type", Value: "exact"}, {Key: "total_cost", Value: 10.0}},
})
require.NoError(t, err)
start, end := now.Add(-time.Hour), now.Add(time.Hour)
got, err := store.SumSpend(ctx, []SpendWindow{
{Scope: ScopeProvider, Subject: "openai-prod", Start: start, End: end},
{Scope: ScopeProvider, Subject: "openai", Start: start, End: end},
{Scope: ScopeProvider, Subject: "anthropic", Start: start, End: end},
})
require.NoError(t, err)
assert.Equal(t, []Spend{
{Total: 2, HasUsage: true}, // case-insensitive, cached row excluded
{Total: 3, HasUsage: true}, // only the row without a provider name
{},
}, got)
})
}
// Rows served by a billable: false provider count against provider budgets
// only; documents without the field predate it and stay billable.
func TestMongoDBStoreSumSpendSkipsNonBillableRowsForConsumerBudgets(t *testing.T) {
mongotest.Run(t, func(t *testing.T, db *mongo.Database) {
ctx := context.Background()
store, err := NewMongoDBStore(ctx, db)
require.NoError(t, err)
now := time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)
_, err = db.Collection("usage").InsertMany(ctx, []any{
bson.D{{Key: "timestamp", Value: now}, {Key: "user_path", Value: "/team/app"}, {Key: "labels", Value: bson.A{"prod"}}, {Key: "provider_name", Value: "openai-prod"}, {Key: "total_cost", Value: 2.0}, {Key: "billable", Value: true}},
bson.D{{Key: "timestamp", Value: now}, {Key: "user_path", Value: "/team/app"}, {Key: "labels", Value: bson.A{"prod"}}, {Key: "provider_name", Value: "byok-openai"}, {Key: "total_cost", Value: 50.0}, {Key: "billable", Value: false}},
// Written before the flag existed.
bson.D{{Key: "timestamp", Value: now}, {Key: "user_path", Value: "/team/app"}, {Key: "labels", Value: bson.A{"prod"}}, {Key: "provider_name", Value: "openai-prod"}, {Key: "total_cost", Value: 3.0}},
})
require.NoError(t, err)
start, end := now.Add(-time.Hour), now.Add(time.Hour)
got, err := store.SumSpend(ctx, []SpendWindow{
{Scope: ScopeUserPath, Subject: "/team", Start: start, End: end},
{Scope: ScopeLabel, Subject: "prod", Start: start, End: end},
{Scope: ScopeProvider, Subject: "byok-openai", Start: start, End: end},
})
require.NoError(t, err)
assert.Equal(t, []Spend{
{Total: 5, HasUsage: true},
{Total: 5, HasUsage: true},
{Total: 50, HasUsage: true},
}, got)
})
}
func TestMongoDBStoreSumSpendMatchesDollarLabelLiterally(t *testing.T) {
mongotest.Run(t, func(t *testing.T, db *mongo.Database) {
ctx := context.Background()
store, err := NewMongoDBStore(ctx, db)
require.NoError(t, err)
now := time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)
_, err = db.Collection("usage").InsertMany(ctx, []any{
bson.D{{Key: "timestamp", Value: now}, {Key: "labels", Value: bson.A{"$labels"}}, {Key: "total_cost", Value: 2.0}},
bson.D{{Key: "timestamp", Value: now}, {Key: "labels", Value: bson.A{"other"}}, {Key: "total_cost", Value: 4.0}},
})
require.NoError(t, err)
got, err := store.SumSpend(ctx, []SpendWindow{
{Scope: ScopeLabel, Subject: "$labels", Start: now.Add(-time.Hour), End: now.Add(time.Hour)},
})
require.NoError(t, err)
assert.Equal(t, []Spend{{Total: 2, HasUsage: true}}, got)
})
}