llmgw/internal/store/postgres_memberships_test.go

package store

import (
	"testing"
	"time"
)

// The membership queries on PostgreSQL: the unions that decide which projects
// a person may use, the path each membership is named by, the revocation that
// follows leaving, a personal key's resolution, and spend with no project.
// Each passed on SQLite first; these are what would catch one that does not
// on the database the cluster runs.
func TestPostgresMemberships(t *testing.T) {
	f := newFixtureOn(t, openPG(t))
	ctx := t.Context()

	if err := f.s.AddMembership(ctx, f.user, ScopeTeam, f.ml, RoleMember, "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	if err := f.s.AddMembership(ctx, f.user, ScopeTeam, f.ml, RoleMember, "t", time.Now()); err != nil {
		t.Errorf("adding again should change nothing: %v", err)
	}
	for proj, want := range map[string]bool{f.gateway: true, f.research: true, f.search: false} {
		if ok, err := f.s.CanUseProject(ctx, f.user, proj); err != nil || ok != want {
			t.Errorf("CanUseProject %s = %v, %v; want %v", proj, ok, err, want)
		}
	}
	if ps, err := f.s.ProjectsFor(ctx, f.user); err != nil || len(ps) != 2 {
		t.Errorf("ProjectsFor = %+v, %v", ps, err)
	}
	ms, err := f.s.MembershipsOf(ctx, f.user)
	if err != nil || len(ms) != 2 || ms[0].Path != "Acme / platform / gateway" || ms[1].Path != "Scania / ml" {
		t.Errorf("MembershipsOf = %+v, %v", ms, err)
	}
	if all, err := f.s.AllMemberships(ctx); err != nil || len(all) != 2 {
		t.Errorf("AllMemberships = %+v, %v", all, err)
	}

	// Leaving gateway revokes the key made there, and leaves a personal one.
	_, personal, _ := NewKey("key-p", "personal", "")
	if err := f.s.CreateUserKey(ctx, f.user, personal); err != nil {
		t.Fatal(err)
	}
	if n, err := f.s.RemoveMembership(ctx, f.user, ScopeProject, f.gateway, time.Now()); err != nil || n != 1 {
		t.Errorf("RemoveMembership revoked %d, %v; want 1", n, err)
	}

	// A personal key resolves, under the personal defaults.
	f.s.SetPersonalDefaults(&Limits{RPM: i(9)})
	defer f.s.SetPersonalDefaults(nil)
	e, err := f.s.Resolve(ctx, personal.ID)
	if err != nil || e.UserID != f.user || e.ProjectID != "" || e.RPM == nil || *e.RPM != 9 {
		t.Errorf("personal key resolved to %+v, %v", e, err)
	}
	if e, err := f.s.ResolveFor(ctx, "session:"+f.user, f.user, f.research); err != nil || e.OrgID != f.scania {
		t.Errorf("ResolveFor in a project = %+v, %v", e, err)
	}

	// Spend with no project is kept, and attributed.
	now := time.Now()
	if err := f.s.WriteUsage(ctx, []Usage{
		{At: now, KeyID: personal.ID, UserID: f.user, CostMicros: 5, Model: "m", ServedModel: "m"},
		{At: now, KeyID: "session:" + f.user, UserID: f.user, ProjectID: f.research, CostMicros: 7, Model: "m", ServedModel: "m"},
	}); err != nil {
		t.Fatal(err)
	}
	rows, err := f.s.SpendSinceByScope(ctx, now.Add(-time.Minute))
	if err != nil {
		t.Fatal(err)
	}
	var total int64
	for _, r := range rows {
		if r.UserID == f.user {
			total += r.Micros
		}
	}
	if total != 12 {
		t.Errorf("restored %d for the person, want 12: %+v", total, rows)
	}
	byUser, err := f.s.UsageByUser(ctx, now.Add(-time.Minute))
	if err != nil || len(byUser) != 1 || byUser[0].UserID != f.user || byUser[0].Requests != 2 {
		t.Errorf("UsageByUser = %+v, %v", byUser, err)
	}
	if keys, err := f.s.AllKeys(ctx); err != nil || len(keys) != 3 {
		t.Errorf("AllKeys = %d rows, %v", len(keys), err)
	}

	// Reports by where it was spent, the chat page's included.
	if series, err := f.s.UsageBucketsBy(ctx, GroupByTeam, now.Add(-time.Minute), time.Minute); err != nil ||
		len(series) != 1 || series[0].Requests != 1 {
		t.Errorf("by team = %+v, %v; the chat request in research should count", series, err)
	}
	if in, err := f.s.UsageBucketsIn(ctx, GroupByUser, f.scania, now.Add(-time.Minute), time.Minute); err != nil ||
		len(in) != 1 || in[0].Requests != 1 {
		t.Errorf("Scania's own use = %+v, %v", in, err)
	}
	if in, err := f.s.UsageBucketsIn(ctx, GroupByUser, f.acme, now.Add(-time.Minute), time.Minute); err != nil || len(in) != 0 {
		t.Errorf("Acme saw use it did not have: %+v, %v", in, err)
	}
	if err := f.s.AddMembership(ctx, f.user, ScopeOrg, f.acme, RoleManager, "t", time.Now()); err != nil {
		t.Fatal(err)
	}
	if ok, err := f.s.ManagesOrg(ctx, f.user, f.acme); err != nil || !ok {
		t.Errorf("ManagesOrg = %v, %v", ok, err)
	}
	if orgs, err := f.s.ManagedOrgs(ctx, f.user); err != nil || len(orgs) != 1 {
		t.Errorf("ManagedOrgs = %+v, %v", orgs, err)
	}
}