Skip to content

Commit 17cc881

Browse files
EmyrkEmyrk
committed
chore: add unit test for testing workspace quota query
1 parent 854044e commit 17cc881

File tree

2 files changed

+106
-2
lines changed

2 files changed

+106
-2
lines changed

‎coderd/database/dbgen/dbgen.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,15 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
288288
if err != nil {
289289
return err
290290
}
291+
292+
if orig.DailyCost > 0 {
293+
err = db.UpdateWorkspaceBuildCostByID(genCtx, database.UpdateWorkspaceBuildCostByIDParams{
294+
ID: buildID,
295+
DailyCost: orig.DailyCost,
296+
})
297+
require.NoError(t, err)
298+
}
299+
291300
build, err = db.GetWorkspaceBuildByID(genCtx, buildID)
292301
if err != nil {
293302
return err

‎enterprise/coderd/workspacequota_test.go

Lines changed: 97 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/coder/coder/v2/coderd/database/dbgen"
2222
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2323
"github.com/coder/coder/v2/coderd/database/dbtime"
24+
"github.com/coder/coder/v2/coderd/rbac"
2425
"github.com/coder/coder/v2/coderd/util/ptr"
2526
"github.com/coder/coder/v2/codersdk"
2627
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
@@ -31,9 +32,13 @@ import (
3132
)
3233

3334
func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, organizationID string, consumed, total int) {
35+
verifyQuotaUser(ctx, t, client, organizationID, codersdk.Me, consumed, total)
36+
}
37+
38+
func verifyQuotaUser(ctx context.Context, t *testing.T, client *codersdk.Client, organizationID string, user string, consumed, total int) {
3439
t.Helper()
3540

36-
got, err := client.WorkspaceQuota(ctx, organizationID, codersdk.Me)
41+
got, err := client.WorkspaceQuota(ctx, organizationID, user)
3742
require.NoError(t, err)
3843
require.EqualValues(t, codersdk.WorkspaceQuota{
3944
Budget: total,
@@ -43,7 +48,7 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, org
4348
// Remove this check when the deprecated endpoint is removed.
4449
// This just makes sure the deprecated endpoint is still working
4550
// as intended. It will only work for the default organization.
46-
deprecatedGot, err := deprecatedQuotaEndpoint(ctx, client, codersdk.Me)
51+
deprecatedGot, err := deprecatedQuotaEndpoint(ctx, client, user)
4752
require.NoError(t, err, "deprecated endpoint")
4853
// Only continue to check if the values differ
4954
if deprecatedGot.Budget != got.Budget || deprecatedGot.CreditsConsumed != got.CreditsConsumed {
@@ -300,6 +305,96 @@ func TestWorkspaceQuota(t *testing.T) {
300305
verifyQuota(ctx, t, owner, first.OrganizationID.String(), 0, 30)
301306
verifyQuota(ctx, t, owner, second.ID.String(), 0, 15)
302307
})
308+
309+
// ManyWorkspaces uses dbfake and dbgen to insert a scenario into the db.
310+
t.Run("ManyWorkspaces", func(t *testing.T) {
311+
t.Parallel()
312+
313+
owner, db, first := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{
314+
LicenseOptions: &coderdenttest.LicenseOptions{
315+
Features: license.Features{
316+
codersdk.FeatureTemplateRBAC: 1,
317+
codersdk.FeatureMultipleOrganizations: 1,
318+
},
319+
},
320+
})
321+
client, _ := coderdtest.CreateAnotherUser(t, owner, first.OrganizationID, rbac.RoleOwner())
322+
323+
// Prepopulate database. Use dbfake as it is quicker and
324+
// easier than the api.
325+
ctx := testutil.Context(t, testutil.WaitLong)
326+
ctx = dbauthz.AsSystemRestricted(ctx)
327+
328+
user := dbgen.User(t, db, database.User{})
329+
noise := dbgen.User(t, db, database.User{})
330+
331+
second := dbfake.Organization(t, db).
332+
Members(user, noise).
333+
EveryoneAllowance(10).
334+
Group(database.Group{
335+
QuotaAllowance: 25,
336+
}, user, noise).
337+
Group(database.Group{
338+
QuotaAllowance: 30,
339+
}, noise).
340+
Do()
341+
342+
third := dbfake.Organization(t, db).
343+
Members(noise).
344+
EveryoneAllowance(7).
345+
Do()
346+
347+
verifyQuotaUser(ctx, t, client, second.Org.ID.String(), user.ID.String(), 0, 35)
348+
verifyQuotaUser(ctx, t, client, second.Org.ID.String(), noise.ID.String(), 0, 65)
349+
350+
// Workspaces owned by the user
351+
consumed := 0
352+
for i := 0; i < 2; i++ {
353+
const cost = 5
354+
dbfake.WorkspaceBuild(t, db,
355+
database.WorkspaceTable{
356+
OwnerID: user.ID,
357+
OrganizationID: second.Org.ID,
358+
}).
359+
Seed(database.WorkspaceBuild{
360+
DailyCost: cost,
361+
}).Do()
362+
consumed += cost
363+
}
364+
365+
// Add some noise
366+
// Workspace by the user in the third org
367+
dbfake.WorkspaceBuild(t, db,
368+
database.WorkspaceTable{
369+
OwnerID: user.ID,
370+
OrganizationID: third.Org.ID,
371+
}).
372+
Seed(database.WorkspaceBuild{
373+
DailyCost: 10,
374+
}).Do()
375+
376+
// Workspace by another user in third org
377+
dbfake.WorkspaceBuild(t, db,
378+
database.WorkspaceTable{
379+
OwnerID: noise.ID,
380+
OrganizationID: third.Org.ID,
381+
}).
382+
Seed(database.WorkspaceBuild{
383+
DailyCost: 10,
384+
}).Do()
385+
386+
// Workspace by another user in second org
387+
dbfake.WorkspaceBuild(t, db,
388+
database.WorkspaceTable{
389+
OwnerID: noise.ID,
390+
OrganizationID: second.Org.ID,
391+
}).
392+
Seed(database.WorkspaceBuild{
393+
DailyCost: 10,
394+
}).Do()
395+
396+
verifyQuotaUser(ctx, t, client, second.Org.ID.String(), user.ID.String(), consumed, 35)
397+
})
303398
}
304399

305400
// nolint:paralleltest,tparallel // Tests must run serially

0 commit comments

Comments
 (0)