llmgw/internal/accounting/people_test.go

package accounting

import (
	"context"
	"errors"
	"log/slog"
	"path/filepath"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// peopleStore has Acme/ml/training with two members, ayse and deniz.
func peopleStore(t *testing.T) (*store.DB, string) {
	t.Helper()
	ctx := context.Background()
	s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "p.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	org, _ := s.CreateOrg(ctx, "Acme")
	team, _ := s.CreateTeam(ctx, org, "ml")
	proj, _ := s.CreateProject(ctx, team, "training")
	for _, id := range []string{"ayse", "deniz"} {
		if err := s.CreateUser(ctx, store.User{ID: id, Email: id + "@x", Name: id,
			PasswordHash: "h", State: store.UserPending, CreatedAt: time.Now()}); err != nil {
			t.Fatal(err)
		}
		if err := s.ApproveUser(ctx, id, proj, "test", time.Now()); err != nil {
			t.Fatal(err)
		}
	}
	return s, proj
}

// keyFor makes a key for a person and resolves it the way authentication does.
func keyFor(t *testing.T, s *store.DB, id, user, project string) auth.Principal {
	t.Helper()
	ctx := context.Background()
	_, k, _ := store.NewKey(id, id, project)
	if err := s.CreateUserKey(ctx, user, k); err != nil {
		t.Fatal(err)
	}
	e, err := s.Resolve(ctx, id)
	if err != nil {
		t.Fatalf("resolving %s: %v", id, err)
	}
	return auth.Principal{KeyID: id, UserID: e.UserID, ProjectID: e.ProjectID, Limits: e}
}

func TestAPersonsLimitIsOnePoolAcrossTheirKeys(t *testing.T) {
	s, proj := peopleStore(t)
	if err := s.PutLimits(context.Background(),
		store.Limits{Scope: store.ScopeUser, ScopeID: "ayse", RPM: intp(2)}); err != nil {
		t.Fatal(err)
	}
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	personal := keyFor(t, s, "k-personal", "ayse", "")
	inProject := keyFor(t, s, "k-project", "ayse", proj)

	// Two requests, one on each key, use up the person's two a minute.
	for _, p := range []auth.Principal{personal, inProject} {
		if _, err := a.Check(p, "free", 1); err != nil {
			t.Fatalf("%s: %v", p.KeyID, err)
		}
	}
	if _, err := a.Check(personal, "free", 1); !errors.Is(err, ErrRateLimited) {
		t.Errorf("a third request across her keys was allowed: %v", err)
	}
	// Deniz is somebody else.
	if _, err := a.Check(keyFor(t, s, "k-deniz", "deniz", proj), "free", 1); err != nil {
		t.Errorf("another person was limited by hers: %v", err)
	}
}

func TestAPerMemberLimitCountsEachMemberApart(t *testing.T) {
	s, proj := peopleStore(t)
	if err := s.PutLimits(context.Background(),
		store.Limits{Scope: store.ScopeProjectMember, ScopeID: proj, RPM: intp(1)}); err != nil {
		t.Fatal(err)
	}
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	ayse := keyFor(t, s, "k-a", "ayse", proj)
	deniz := keyFor(t, s, "k-d", "deniz", proj)

	if _, err := a.Check(ayse, "free", 1); err != nil {
		t.Fatal(err)
	}
	if _, err := a.Check(ayse, "free", 1); !errors.Is(err, ErrRateLimited) {
		t.Errorf("her second request passed a one-a-minute per-member limit: %v", err)
	}
	// Her share is hers: Deniz still has his.
	if _, err := a.Check(deniz, "free", 1); err != nil {
		t.Errorf("one member's use counted against another: %v", err)
	}
	// And outside the project it does not apply at all.
	if _, err := a.Check(keyFor(t, s, "k-a-own", "ayse", ""), "free", 1); err != nil {
		t.Errorf("a project's per-member limit held her personal key: %v", err)
	}
}

func TestAPersonsBudgetSurvivesARestart(t *testing.T) {
	s, _ := peopleStore(t)
	ctx := context.Background()
	if err := s.PutLimits(ctx, store.Limits{Scope: store.ScopeUser, ScopeID: "ayse",
		Budget: i64p(1000), BudgetPeriod: store.PeriodMonth}); err != nil {
		t.Fatal(err)
	}
	// Spent before the restart, personally and on the chat page.
	if err := s.WriteUsage(ctx, []store.Usage{
		{At: time.Now(), KeyID: "k-old", UserID: "ayse", CostMicros: 600, Model: "paid", ServedModel: "paid"},
		{At: time.Now(), KeyID: "session:ayse", UserID: "ayse", CostMicros: 500, Model: "paid", ServedModel: "paid"},
	}); err != nil {
		t.Fatal(err)
	}
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	if err := a.RestoreBudgets(ctx, s, store.PeriodMonth); err != nil {
		t.Fatal(err)
	}
	if _, err := a.Check(keyFor(t, s, "k-new", "ayse", ""), "paid", 1); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("a restart forgot what she had spent: %v", err)
	}
}