package dbfake import ( "testing" "github.com/google/uuid" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/testutil" ) type OrganizationBuilder struct { t *testing.T db database.Store seed database.Organization delete bool allUsersAllowance int32 members []uuid.UUID groups map[database.Group][]uuid.UUID } func Organization(t *testing.T, db database.Store) OrganizationBuilder { return OrganizationBuilder{ t: t, db: db, members: []uuid.UUID{}, groups: make(map[database.Group][]uuid.UUID), } } type OrganizationResponse struct { Org database.Organization AllUsersGroup database.Group Members []database.OrganizationMember Groups []database.Group } func (b OrganizationBuilder) EveryoneAllowance(allowance int) OrganizationBuilder { //nolint: revive // returns modified struct // #nosec G115 - Safe conversion as allowance is expected to be within int32 range b.allUsersAllowance = int32(allowance) return b } func (b OrganizationBuilder) Deleted(deleted bool) OrganizationBuilder { //nolint: revive // returns modified struct b.delete = deleted return b } func (b OrganizationBuilder) Seed(seed database.Organization) OrganizationBuilder { //nolint: revive // returns modified struct b.seed = seed return b } func (b OrganizationBuilder) Members(users ...database.User) OrganizationBuilder { for _, u := range users { //nolint: revive // returns modified struct b.members = append(b.members, u.ID) } return b } func (b OrganizationBuilder) Group(seed database.Group, members ...database.User) OrganizationBuilder { //nolint: revive // returns modified struct b.groups[seed] = []uuid.UUID{} for _, u := range members { //nolint: revive // returns modified struct b.groups[seed] = append(b.groups[seed], u.ID) } return b } func (b OrganizationBuilder) Do() OrganizationResponse { org := dbgen.Organization(b.t, b.db, b.seed) ctx := testutil.Context(b.t, testutil.WaitShort) //nolint:gocritic // builder code needs perms ctx = dbauthz.AsSystemRestricted(ctx) everyone, err := b.db.InsertAllUsersGroup(ctx, org.ID) require.NoError(b.t, err) if b.allUsersAllowance > 0 { everyone, err = b.db.UpdateGroupByID(ctx, database.UpdateGroupByIDParams{ Name: everyone.Name, DisplayName: everyone.DisplayName, AvatarURL: everyone.AvatarURL, QuotaAllowance: b.allUsersAllowance, ID: everyone.ID, }) require.NoError(b.t, err) } members := make([]database.OrganizationMember, 0) if len(b.members) > 0 { for _, u := range b.members { newMem := dbgen.OrganizationMember(b.t, b.db, database.OrganizationMember{ UserID: u, OrganizationID: org.ID, CreatedAt: dbtime.Now(), UpdatedAt: dbtime.Now(), Roles: nil, }) members = append(members, newMem) } } groups := make([]database.Group, 0) if len(b.groups) > 0 { for g, users := range b.groups { g.OrganizationID = org.ID group := dbgen.Group(b.t, b.db, g) groups = append(groups, group) for _, u := range users { dbgen.GroupMember(b.t, b.db, database.GroupMemberTable{ UserID: u, GroupID: group.ID, }) } } } if b.delete { now := dbtime.Now() err = b.db.UpdateOrganizationDeletedByID(ctx, database.UpdateOrganizationDeletedByIDParams{ UpdatedAt: now, ID: org.ID, }) require.NoError(b.t, err) org.Deleted = true org.UpdatedAt = now } return OrganizationResponse{ Org: org, AllUsersGroup: everyone, Members: members, Groups: groups, } }