mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add GET /groups/{group}/members/ai/spend (#27130)
## Description
Adds `GET /api/v2/groups/{group}/members/ai/spend?user_ids=...` (also available org-scoped at `/api/v2/organizations/{org}/groups/{groupName}/members/ai/spend`) to return per-member AI spend attributed to a group, along with each member's effective budget group and the applied spend limit when the queried group is their effective budget source.
In the UI, this endpoint is used alongside the existing `/api/v2/groups/{group}/members` endpoint. AI spend data is kept separate from that endpoint so that:
- Different concepts stay on different endpoints: identity (group members) vs. cost control (spend). Cost control is an additional feature layered on top of groups/orgs.
- Callers that don't need spend information don't pay for its computation.
UI flow:
1. Request `/api/v2/groups/{group}/members` → returns the group's members.
2. Request `/api/v2/groups/{group}/members/ai/spend?user_ids=...` with the IDs from step 1.
**Note:** Only current members of the queried group are returned. `spend_limit_micros` and `limit_source` are populated only when the queried group is the member's effective budget source (its own limit or a user override). `effective_group_id` is null when the member's budget resolves to a group in another organization, since an organization is treated as a tenant boundary.
<img width="2880" height="1904" alt="image" src="https://github.com/user-attachments/assets/33ed395d-d1a3-4b46-bb04-c8d3f41c8886" />
## Changes
- Add `codersdk.GroupMembersAISpend` and `GroupMemberAISpend` types, reusing the shared `AISpendPeriodWindow`.
- Add `GetGroupMembersAISpend` SQL query with a dbauthz per-row filter that mirrors `GET /api/v2/groups/{group}/members`.
- Add handler and routes under `/groups/{group}/members/ai/spend` (and the org-scoped alias) with a required `user_ids` query param (cap 100). Callers with more than 100 members are expected to batch across multiple requests.
- Add codersdk client method.
- Tests: dbauthz, raw SQL, endpoint, and role-access.
Closes https://linear.app/codercom/issue/AIGOV-471/backend-group-members-endpoint-with-members-spend
> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
@@ -13106,6 +13106,592 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetGroupMembersAISpend(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Use fixed dates to keep the test deterministic.
|
||||
monthStart := time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC)
|
||||
now := monthStart.AddDate(0, 0, 14) // 2024-06-15
|
||||
prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
groupLimit int64
|
||||
overrideLimit int64
|
||||
spend int64
|
||||
wantEffectiveGroup bool
|
||||
wantLimit sql.NullInt64
|
||||
wantSource sql.NullString
|
||||
wantSpend int64
|
||||
}{
|
||||
{
|
||||
name: "NoBudgetNoSpend",
|
||||
wantEffectiveGroup: false,
|
||||
wantLimit: sql.NullInt64{},
|
||||
wantSource: sql.NullString{},
|
||||
wantSpend: 0,
|
||||
},
|
||||
{
|
||||
name: "GroupBudget",
|
||||
groupLimit: 1_000_000,
|
||||
wantEffectiveGroup: true,
|
||||
wantLimit: sql.NullInt64{Int64: 1_000_000, Valid: true},
|
||||
wantSource: sql.NullString{String: "group", Valid: true},
|
||||
wantSpend: 0,
|
||||
},
|
||||
{
|
||||
name: "OverrideBudget",
|
||||
overrideLimit: 500_000,
|
||||
wantEffectiveGroup: true,
|
||||
wantLimit: sql.NullInt64{Int64: 500_000, Valid: true},
|
||||
wantSource: sql.NullString{String: "user_override", Valid: true},
|
||||
wantSpend: 0,
|
||||
},
|
||||
{
|
||||
name: "NoBudgetWithSpend",
|
||||
spend: 250,
|
||||
wantEffectiveGroup: false,
|
||||
wantLimit: sql.NullInt64{},
|
||||
wantSource: sql.NullString{},
|
||||
wantSpend: 250,
|
||||
},
|
||||
{
|
||||
name: "BudgetWithSpend",
|
||||
groupLimit: 1_000_000,
|
||||
spend: 250,
|
||||
wantEffectiveGroup: true,
|
||||
wantLimit: sql.NullInt64{Int64: 1_000_000, Valid: true},
|
||||
wantSource: sql.NullString{String: "group", Valid: true},
|
||||
wantSpend: 250,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a member of the queried group, optionally with a group
|
||||
// budget, a user override, and seeded spend.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
if tt.groupLimit > 0 {
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: group.ID,
|
||||
SpendLimitMicros: tt.groupLimit,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
if tt.overrideLimit > 0 {
|
||||
_, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
|
||||
UserID: user.ID,
|
||||
GroupID: group.ID,
|
||||
SpendLimitMicros: tt.overrideLimit,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
if tt.spend > 0 {
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: tt.spend,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// When: querying spend for the user.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: one row with the expected effective group, limit, source, and spend.
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, user.ID, got[0].UserID)
|
||||
require.Equal(t, org.ID, got[0].OrganizationID)
|
||||
if tt.wantEffectiveGroup {
|
||||
require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID)
|
||||
} else {
|
||||
require.False(t, got[0].EffectiveGroupID.Valid, "expected no effective group")
|
||||
}
|
||||
require.Equal(t, tt.wantLimit, got[0].SpendLimitMicros)
|
||||
require.Equal(t, tt.wantSource, got[0].LimitSource)
|
||||
require.Equal(t, tt.wantSpend, got[0].GroupSpendMicros)
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("MultipleMembers", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: two group members with different spend attributed to the group.
|
||||
userA := dbgen.User(t, db, database.User{})
|
||||
userB := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userA.ID, OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userB.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userA.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userB.ID})
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: userA.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 100,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: userB.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying spend for both users.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{userA.ID, userB.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: both users are returned with their own aggregate spend.
|
||||
require.Len(t, got, 2)
|
||||
byID := make(map[uuid.UUID]database.GetGroupMembersAISpendRow, len(got))
|
||||
for _, row := range got {
|
||||
byID[row.UserID] = row
|
||||
}
|
||||
require.Equal(t, int64(100), byID[userA.ID].GroupSpendMicros)
|
||||
require.Equal(t, int64(250), byID[userB.ID].GroupSpendMicros)
|
||||
for _, row := range got {
|
||||
require.False(t, row.EffectiveGroupID.Valid)
|
||||
require.False(t, row.SpendLimitMicros.Valid)
|
||||
require.False(t, row.LimitSource.Valid)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("AggregatesSpendAcrossDays", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a group member with spend on multiple days in the period.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
days := []time.Time{monthStart, monthStart.AddDate(0, 0, 7), monthStart.AddDate(0, 0, 14)}
|
||||
for i, day := range days {
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: int64((i + 1) * 100),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// When: querying the group's member spend.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: per-user spend is summed across all days in the period.
|
||||
require.Len(t, got, 1)
|
||||
require.False(t, got[0].EffectiveGroupID.Valid)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(600), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("OverrideWins", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a member of the queried group who is also in two other
|
||||
// budgeted groups, one being the natural highest-limit and the other
|
||||
// set via a user override.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
overrideTarget := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
highestLimit := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: overrideTarget.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: highestLimit.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: overrideTarget.ID,
|
||||
SpendLimitMicros: 1_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: highestLimit.ID,
|
||||
SpendLimitMicros: 5_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
|
||||
UserID: user.ID,
|
||||
GroupID: overrideTarget.ID,
|
||||
SpendLimitMicros: 500_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying spend for the user.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: queried.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: the override target wins over the highest-limit group.
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, uuid.NullUUID{UUID: overrideTarget.ID, Valid: true}, got[0].EffectiveGroupID)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(0), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("EqualBudgetTieBreak", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a member of the queried group who is in two same-org groups
|
||||
// with identical spend limits.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "aaa-tie-group"})
|
||||
groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "bbb-tie-group"})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupA.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: groupA.ID,
|
||||
SpendLimitMicros: 1_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: groupB.ID,
|
||||
SpendLimitMicros: 1_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying spend for the user.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: queried.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: the tie is broken by group name ascending, so groupA wins.
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, uuid.NullUUID{UUID: groupA.ID, Valid: true}, got[0].EffectiveGroupID)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(0), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("EveryoneGroupCounts", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a member of the queried group whose only budgeted group is
|
||||
// the org's implicit Everyone group.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
|
||||
// The Everyone group has ID equal to the organization ID and must be
|
||||
// inserted explicitly for this test's FK constraint on group_ai_budgets.
|
||||
//nolint:gocritic // Requires system context.
|
||||
_, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: org.ID,
|
||||
SpendLimitMicros: 1_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying spend for the user.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: queried.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: effective_group_id resolves to the Everyone group.
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(0), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("SpendWithDifferentEffectiveGroup", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a group member with spend attributed to the queried group,
|
||||
// whose current effective group is a different same-org group.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
other := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: other.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: other.ID,
|
||||
SpendLimitMicros: 1_000_000,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying the queried group's spend for the user.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: queried.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: queried-group spend is returned, effective_group_id is the
|
||||
// other group, and the limit and source are null because the queried
|
||||
// group is not the effective source.
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, uuid.NullUUID{UUID: other.ID, Valid: true}, got[0].EffectiveGroupID)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(250), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("ExcludesOtherGroupSpend", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a group member with spend attributed to a different group in the same org.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
otherGroup := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: otherGroup.ID, Day: now, CostMicros: 500,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying spend for the queried group.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: spend attributed to the other group is not counted.
|
||||
require.Len(t, got, 1)
|
||||
require.False(t, got[0].EffectiveGroupID.Valid)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(0), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("ExcludesNonMembers", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a member of the queried group and a non-member in the same org.
|
||||
member := dbgen.User(t, db, database.User{})
|
||||
nonMember := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: member.ID, OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: nonMember.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID})
|
||||
|
||||
// When: querying with both user IDs.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{member.ID, nonMember.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: only the group member is returned.
|
||||
require.Len(t, got, 1)
|
||||
require.Equal(t, member.ID, got[0].UserID)
|
||||
require.False(t, got[0].EffectiveGroupID.Valid)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(0), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("HidesFormerMember", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a user who has historical spend attributed to the queried
|
||||
// group but is not currently a member of it.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 500,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying the queried group's spend for the user.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: queried.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: the user is filtered out and their historical spend is not returned.
|
||||
require.Empty(t, got)
|
||||
})
|
||||
|
||||
t.Run("CrossOrgEffectiveGroupMasked", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a member of the queried group whose highest-limit budget
|
||||
// group is in a different org.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
otherOrg := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
otherOrgGroup := dbgen.Group(t, db, database.Group{OrganizationID: otherOrg.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: otherOrgGroup.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
|
||||
GroupID: otherOrgGroup.ID,
|
||||
SpendLimitMicros: 9_999_999,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
// Seed spend attributed to the queried group so we can assert it is
|
||||
// still returned even when the effective group is masked.
|
||||
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying spend for the user in the queried group's org.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: effective_group_id is masked to NULL, the highest-limit group is cross-org.
|
||||
// The queried-group spend is still returned.
|
||||
require.Len(t, got, 1)
|
||||
require.False(t, got[0].EffectiveGroupID.Valid, "cross-org effective group must be masked")
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(250), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("ExcludesSpendBeforePeriodStart", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a group member with spend both in the prior period and in the current period.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying since monthStart.
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: monthStart,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: only current-period spend is aggregated.
|
||||
require.Len(t, got, 1)
|
||||
require.False(t, got[0].EffectiveGroupID.Valid)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(25), got[0].GroupSpendMicros)
|
||||
})
|
||||
|
||||
t.Run("NormalizesNonUTCPeriodStart", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Given: a group member with spend on the prior UTC day and on the first day of the current UTC month.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: querying with a non-UTC period_start that normalizes to June 1 UTC.
|
||||
// 2024-05-31 23:00 in UTC-5 is 2024-06-01 04:00 UTC.
|
||||
localLate := time.Date(2024, 5, 31, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600))
|
||||
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
|
||||
GroupID: group.ID,
|
||||
UserIds: []uuid.UUID{user.ID},
|
||||
PeriodStart: localLate,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: the prior UTC day's spend is excluded from the aggregate.
|
||||
require.Len(t, got, 1)
|
||||
require.False(t, got[0].EffectiveGroupID.Valid)
|
||||
require.False(t, got[0].SpendLimitMicros.Valid)
|
||||
require.False(t, got[0].LimitSource.Valid)
|
||||
require.Equal(t, int64(25), got[0].GroupSpendMicros,
|
||||
"sum must exclude prevMonthLastDay row after normalization")
|
||||
})
|
||||
}
|
||||
|
||||
func TestChatPinOrderQueries(t *testing.T) {
|
||||
t.Parallel()
|
||||
if testing.Short() {
|
||||
|
||||
Reference in New Issue
Block a user