llmgw/internal/store/memberships_test.go

package store

import (
	"errors"
	"testing"
	"time"
)

func TestApprovalNeedNotPutAnybodyAnywhere(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	_ = f.s.CreateUser(ctx, User{ID: "usr-2", Email: "b@x", Name: "b", PasswordHash: "h",
		State: UserPending, CreatedAt: time.Now()})
	if err := f.s.ApproveUser(ctx, "usr-2", "", "test", time.Now()); err != nil {
		t.Fatalf("approving with no project: %v", err)
	}
	ms, _ := f.s.MembershipsOf(ctx, "usr-2")
	if len(ms) != 0 {
		t.Errorf("memberships = %+v", ms)
	}
}

func TestATeamMemberUsesEveryProjectInTheTeam(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	if err := f.s.AddMembership(ctx, f.user, ScopeTeam, f.ml, RoleMember, "test", time.Now()); err != nil {
		t.Fatal(err)
	}
	for proj, want := range map[string]bool{f.gateway: true, f.research: true, f.search: false} {
		ok, err := f.s.CanUseProject(ctx, f.user, proj)
		if err != nil || ok != want {
			t.Errorf("%s: %v %v, want %v", proj, ok, err, want)
		}
	}
	projects, _ := f.s.ProjectsFor(ctx, f.user)
	if len(projects) != 2 {
		t.Errorf("projects = %+v", projects)
	}
	ms, _ := f.s.MembershipsOf(ctx, f.user)
	var paths []string
	for _, m := range ms {
		paths = append(paths, m.Path)
	}
	if len(paths) != 2 || paths[0] != "Acme / platform / gateway" || paths[1] != "Scania / ml" {
		t.Errorf("paths = %v", paths)
	}
}

func TestAKeyCanOnlyBeMadeWhereThePersonIsAMember(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	_, k, _ := NewKey("key-x", "x", f.search)
	if err := f.s.CreateUserKey(ctx, f.user, k); !errors.Is(err, ErrNotMember) {
		t.Errorf("a key in a project the person is not in: %v", err)
	}
	_, k, _ = NewKey("key-p", "personal", "")
	if err := f.s.CreateUserKey(ctx, f.user, k); err != nil {
		t.Fatalf("a personal key: %v", err)
	}
	got, err := f.s.KeyByHash(ctx, k.Hash)
	if err != nil || got.ProjectID != "" || got.UserID != f.user {
		t.Errorf("personal key read back as %+v, %v", got, err)
	}
}

func TestLeavingAProjectRevokesTheKeysMadeThereOnly(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	_, personal, _ := NewKey("key-p", "personal", "")
	if err := f.s.CreateUserKey(ctx, f.user, personal); err != nil {
		t.Fatal(err)
	}
	revoked, err := f.s.RemoveMembership(ctx, f.user, ScopeProject, f.gateway, time.Now())
	if err != nil {
		t.Fatal(err)
	}
	if revoked != 1 {
		t.Errorf("revoked %d keys, want the one made in gateway", revoked)
	}
	keys, _ := f.s.KeysForUser(ctx, f.user)
	for _, k := range keys {
		if k.ID == personal.ID && !k.RevokedAt.IsZero() {
			t.Error("leaving a project revoked a personal key")
		}
	}
	gw, _ := f.s.Keys(ctx, f.gateway)
	for _, k := range gw {
		if k.ID == f.operatorKey && !k.RevokedAt.IsZero() {
			t.Error("an operator's key went with the person")
		}
	}
}

func TestLeavingKeepsKeysReachedAnotherWay(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	// Also a member of the whole organization, so gateway stays usable.
	if err := f.s.AddMembership(ctx, f.user, ScopeOrg, f.acme, RoleMember, "", time.Now()); err != nil {
		t.Fatal(err)
	}
	revoked, err := f.s.RemoveMembership(ctx, f.user, ScopeProject, f.gateway, time.Now())
	if err != nil || revoked != 0 {
		t.Errorf("revoked %d, %v; the organization still reaches gateway", revoked, err)
	}
}

func i(v int) *int { return &v }

func TestAPersonalKeyIsBoundedByThePerson(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	_, k, _ := NewKey("key-p", "personal", "")
	if err := f.s.CreateUserKey(ctx, f.user, k); err != nil {
		t.Fatal(err)
	}

	// Nothing written for the person: the personal defaults apply, in the
	// person's own pool.
	f.s.SetPersonalDefaults(&Limits{RPM: i(10), TPM: i(5000)})
	e, err := f.s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatalf("resolving a personal key: %v", err)
	}
	if e.ProjectID != "" || e.UserID != f.user || e.RPM == nil || *e.RPM != 10 ||
		e.RPMOwner != (LimitOwner{ScopeUser, f.user}) {
		t.Errorf("personal key resolved to %+v", e)
	}

	// Written for the person: that wins over the defaults.
	if err := f.s.PutLimits(ctx, Limits{Scope: ScopeUser, ScopeID: f.user, RPM: i(3)}); err != nil {
		t.Fatal(err)
	}
	e, _ = f.s.Resolve(ctx, k.ID)
	if *e.RPM != 3 || e.TPM != nil {
		t.Errorf("the person's own limit should replace the defaults: %+v", e)
	}
}

func TestAProjectKeyAnswersToThePersonAndEachMemberLimit(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	_, k, _ := NewKey("key-g", "gw", f.gateway)
	if err := f.s.CreateUserKey(ctx, f.user, k); err != nil {
		t.Fatal(err)
	}
	must := func(l Limits) {
		t.Helper()
		if err := f.s.PutLimits(ctx, l); err != nil {
			t.Fatal(err)
		}
	}
	must(Limits{Scope: ScopeProject, ScopeID: f.gateway, RPM: i(100)})
	must(Limits{Scope: ScopeProjectMember, ScopeID: f.gateway, RPM: i(20)})
	must(Limits{Scope: ScopeUser, ScopeID: f.user, TPM: i(700)})

	e, err := f.s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if *e.RPM != 20 || e.RPMOwner != MemberOwner(ScopeProjectMember, f.gateway, f.user) {
		t.Errorf("the per-member limit should hold this person to 20: %+v", e)
	}
	if *e.TPM != 700 || e.TPMOwner != (LimitOwner{ScopeUser, f.user}) {
		t.Errorf("the person's own limit should apply in a project too: %+v", e)
	}
	// Every pool is kept for enforcement, not just the tightest.
	var scopes []Scope
	for _, l := range e.Chain {
		scopes = append(scopes, l.Scope)
	}
	if len(scopes) < 3 {
		t.Errorf("chain = %v", scopes)
	}

	// An operator's key in the same project has no member and no person.
	e, _ = f.s.Resolve(ctx, f.operatorKey)
	if *e.RPM != 100 {
		t.Errorf("an ownerless key is bounded by the project alone: %+v", e)
	}
}

func TestSpendIsRestoredForPeopleAndPersonalKeys(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	now := time.Now()
	if err := f.s.WriteUsage(ctx, []Usage{
		{At: now, KeyID: "key-p", UserID: f.user, CostMicros: 5, Model: "m", ServedModel: "m"},
		{At: now, KeyID: "key-u", UserID: f.user, ProjectID: f.gateway, CostMicros: 7, Model: "m", ServedModel: "m"},
		{At: now, KeyID: "session:" + f.user, UserID: f.user, CostMicros: 11, Model: "m", ServedModel: "m"},
	}); err != nil {
		t.Fatal(err)
	}
	rows, err := f.s.SpendSinceByScope(ctx, now.Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	var personal, project, total int64
	for _, r := range rows {
		if r.UserID != f.user {
			t.Errorf("row not attributed to the person: %+v", r)
		}
		total += r.Micros
		if r.ProjectID == "" {
			personal += r.Micros
		} else if r.OrgID == f.acme {
			project += r.Micros
		}
	}
	if total != 23 || personal != 16 || project != 7 {
		t.Errorf("total %d personal %d project %d", total, personal, project)
	}
}

func TestAManagerIsWhatTheMembershipSays(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	if err := f.s.AddMembership(ctx, f.user, ScopeOrg, f.acme, RoleManager, "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	if ok, _ := f.s.ManagesOrg(ctx, f.user, f.acme); !ok {
		t.Error("a manager of Acme is not recognised as one")
	}
	if ok, _ := f.s.ManagesOrg(ctx, f.user, f.scania); ok {
		t.Error("managing Acme made them a manager of Scania")
	}
	orgs, _ := f.s.ManagedOrgs(ctx, f.user)
	if len(orgs) != 1 || orgs[0].Name != "Acme" {
		t.Errorf("ManagedOrgs = %+v", orgs)
	}
	// Back to member: no longer a manager.
	if err := f.s.AddMembership(ctx, f.user, ScopeOrg, f.acme, RoleMember, "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	if ok, _ := f.s.ManagesOrg(ctx, f.user, f.acme); ok {
		t.Error("changing the role back did not take")
	}
}

func TestAnOrganizationSeesOnlyWhatWasSpentInIt(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	now := time.Now()
	if err := f.s.WriteUsage(ctx, []Usage{
		{At: now, KeyID: "key-u", UserID: f.user, ProjectID: f.gateway, CostMicros: 1, Model: "m", ServedModel: "m"},
		{At: now, KeyID: "session:" + f.user, UserID: f.user, ProjectID: f.research, CostMicros: 1, Model: "m", ServedModel: "m"},
		{At: now, KeyID: "key-p", UserID: f.user, CostMicros: 1, Model: "m", ServedModel: "m"},
	}); err != nil {
		t.Fatal(err)
	}
	rows, err := f.s.UsageBucketsIn(ctx, GroupByUser, f.acme, now.Add(-time.Hour), time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 || rows[0].Requests != 1 {
		t.Errorf("Acme should see only the request made in Acme, got %+v", rows)
	}
}

func TestTheChatPagesUseIsReportedWhereItWasSpent(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	now := time.Now()
	// The chat page's requests have no key row; joined through api_keys they
	// vanished from every report but the per-person one.
	if err := f.s.WriteUsage(ctx, []Usage{
		{At: now, KeyID: "session:" + f.user, UserID: f.user, ProjectID: f.gateway,
			CostMicros: 4, Model: "m", ServedModel: "m"},
	}); err != nil {
		t.Fatal(err)
	}
	series, err := f.s.UsageBucketsBy(ctx, GroupByTeam, now.Add(-time.Hour), time.Hour)
	if err != nil || len(series) != 1 || series[0].Requests != 1 {
		t.Errorf("by team: %+v, %v", series, err)
	}
	scopes, err := f.s.UsageByScope(ctx, ScopeProject, now.Add(-time.Hour))
	var got int64
	for _, sc := range scopes {
		if sc.ID == f.gateway {
			got = sc.CostMicros
		}
	}
	if err != nil || got != 4 {
		t.Errorf("gateway's spend on the quota page: %d, %v", got, err)
	}
}

func TestASuspendedPersonsKeysStopWorkingAndComeBackWithThem(t *testing.T) {
	f := newFixture(t)
	ctx := t.Context()
	_, k, _ := NewKey("key-s", "personal", "")
	if err := f.s.CreateUserKey(ctx, f.user, k); err != nil {
		t.Fatal(err)
	}
	if _, err := f.s.Resolve(ctx, k.ID); err != nil {
		t.Fatalf("an active person's key: %v", err)
	}
	if err := f.s.SetUserState(ctx, f.user, UserSuspended); err != nil {
		t.Fatal(err)
	}
	if _, err := f.s.Resolve(ctx, k.ID); !errors.Is(err, ErrOwnerInactive) {
		t.Errorf("a suspended person's key resolved: %v", err)
	}
	if err := f.s.SetUserState(ctx, f.user, UserActive); err != nil {
		t.Fatal(err)
	}
	if _, err := f.s.Resolve(ctx, k.ID); err != nil {
		t.Errorf("reinstated, the key still refused: %v", err)
	}
}