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) } }