Skip to content

Commit 3abc329

Browse files
EmyrkEmyrk
committed
update unit tests for quotas
1 parent 7f0eca5 commit 3abc329

File tree

4 files changed

+75
-30
lines changed

4 files changed

+75
-30
lines changed

‎codersdk/workspaces.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -572,8 +572,8 @@ type WorkspaceQuota struct {
572572
Budget int `json:"budget"`
573573
}
574574

575-
func (c *Client) WorkspaceQuota(ctx context.Context, userID string) (WorkspaceQuota, error) {
576-
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil)
575+
func (c *Client) WorkspaceQuota(ctx context.Context, organizationID string, userID string) (WorkspaceQuota, error) {
576+
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/organizations/%s/members/%s/workspace-quota", organizationID, userID), nil)
577577
if err != nil {
578578
return WorkspaceQuota{}, err
579579
}

‎enterprise/coderd/coderd.go

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,19 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
274274
r.Delete("/organizations/{organization}/members/roles/{roleName}", api.deleteOrgRole)
275275
})
276276

277+
r.Group(func(r chi.Router) {
278+
r.Use(
279+
apiKeyMiddleware,
280+
httpmw.ExtractOrganizationParam(api.Database),
281+
// Intentionally using ExtractUser instead of ExtractMember.
282+
// It is possible for a member to be removed from an org, in which
283+
// case their orphaned workspaces still exist. We only need
284+
// the user_id for the query.
285+
httpmw.ExtractUserParam(api.Database),
286+
)
287+
r.Get("/organizations/{organization}/members/{user}/workspace-quota", api.workspaceQuota)
288+
})
289+
277290
r.Route("/organizations/{organization}/groups", func(r chi.Router) {
278291
r.Use(
279292
apiKeyMiddleware,
@@ -364,7 +377,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
364377
)
365378
r.Route("/{user}", func(r chi.Router) {
366379
r.Use(httpmw.ExtractUserParam(options.Database))
367-
r.Get("/", api.workspaceQuota)
380+
r.Get("/", api.workspaceQuotaByUser)
368381
})
369382
})
370383
r.Route("/appearance", func(r chi.Router) {

‎enterprise/coderd/workspacequota.go

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -134,10 +134,11 @@ func (api *API) workspaceQuotaByUser(rw http.ResponseWriter, r *http.Request) {
134134
httpapi.InternalServerError(rw, err)
135135
return
136136
}
137-
chi.RouteContext(r.Context()).URLParams.Add("organization", defaultOrg.ID.String())
138137

139-
// defer to the new endpoint
140-
api.workspaceQuota(rw, r)
138+
// defer to the new endpoint using default org as the organization
139+
chi.RouteContext(r.Context()).URLParams.Add("organization", defaultOrg.ID.String())
140+
mw := httpmw.ExtractOrganizationParam(api.Database)
141+
mw(http.HandlerFunc(api.workspaceQuota)).ServeHTTP(rw, r)
141142
}
142143

143144
// @Summary Get workspace quota by user
@@ -151,11 +152,9 @@ func (api *API) workspaceQuotaByUser(rw http.ResponseWriter, r *http.Request) {
151152
func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
152153
var (
153154
organization = httpmw.OrganizationParam(r)
154-
member = httpmw.OrganizationMemberParam(r)
155+
user = httpmw.UserParam(r)
155156
)
156157

157-
//defaultOrg, _ := api.Database.GetDefaultOrganization(r.Context())
158-
159158
api.entitlementsMu.RLock()
160159
licensed := api.entitlements.Features[codersdk.FeatureTemplateRBAC].Enabled
161160
api.entitlementsMu.RUnlock()
@@ -165,7 +164,7 @@ func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
165164
if licensed {
166165
var err error
167166
quotaAllowance, err = api.Database.GetQuotaAllowanceForUser(r.Context(), database.GetQuotaAllowanceForUserParams{
168-
UserID: member.UserID,
167+
UserID: user.ID,
169168
OrganizationID: organization.ID,
170169
})
171170
if err != nil {
@@ -178,7 +177,7 @@ func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
178177
}
179178

180179
quotaConsumed, err := api.Database.GetQuotaConsumedForUser(r.Context(), database.GetQuotaConsumedForUserParams{
181-
OwnerID: member.UserID,
180+
OwnerID: user.ID,
182181
OrganizationID: organization.ID,
183182
})
184183
if err != nil {

‎enterprise/coderd/workspacequota_test.go

Lines changed: 52 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ package coderd_test
22

33
import (
44
"context"
5+
"encoding/json"
6+
"fmt"
7+
"net/http"
58
"sync"
69
"testing"
710

@@ -20,15 +23,31 @@ import (
2023
"github.com/coder/coder/v2/testutil"
2124
)
2225

23-
func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, consumed, total int) {
26+
func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, organizationID string, consumed, total int) {
2427
t.Helper()
2528

26-
got, err := client.WorkspaceQuota(ctx, codersdk.Me)
29+
got, err := client.WorkspaceQuota(ctx, organizationID, codersdk.Me)
2730
require.NoError(t, err)
2831
require.EqualValues(t, codersdk.WorkspaceQuota{
2932
Budget: total,
3033
CreditsConsumed: consumed,
3134
}, got)
35+
36+
// Remove this check when the deprecated endpoint is removed.
37+
// This just makes sure the deprecated endpoint is still working
38+
// as intended. It will only work for the default organization.
39+
deprecatedGot, err := deprecatedQuotaEndpoint(ctx, client, codersdk.Me)
40+
require.NoError(t, err, "deprecated endpoint")
41+
// Only continue to check if the values differ
42+
if deprecatedGot.Budget != got.Budget || deprecatedGot.CreditsConsumed != got.CreditsConsumed {
43+
org, err := client.OrganizationByName(ctx, organizationID)
44+
if err != nil {
45+
return
46+
}
47+
if org.IsDefault {
48+
require.Equal(t, got, deprecatedGot)
49+
}
50+
}
3251
}
3352

3453
func TestWorkspaceQuota(t *testing.T) {
@@ -52,14 +71,14 @@ func TestWorkspaceQuota(t *testing.T) {
5271
})
5372
coderdtest.NewProvisionerDaemon(t, api.AGPL)
5473

55-
verifyQuota(ctx, t, client, 0, 0)
74+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0)
5675

5776
// Patch the 'Everyone' group to verify its quota allowance is being accounted for.
5877
_, err := client.PatchGroup(ctx, user.OrganizationID, codersdk.PatchGroupRequest{
5978
QuotaAllowance: ptr.Ref(1),
6079
})
6180
require.NoError(t, err)
62-
verifyQuota(ctx, t, client, 0, 1)
81+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 1)
6382

6483
// Add user to two groups, granting them a total budget of 4.
6584
group1, err := client.CreateGroup(ctx, user.OrganizationID, codersdk.CreateGroupRequest{
@@ -84,7 +103,7 @@ func TestWorkspaceQuota(t *testing.T) {
84103
})
85104
require.NoError(t, err)
86105

87-
verifyQuota(ctx, t, client, 0, 4)
106+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 4)
88107

89108
authToken := uuid.NewString()
90109
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
@@ -123,14 +142,14 @@ func TestWorkspaceQuota(t *testing.T) {
123142
}()
124143
}
125144
wg.Wait()
126-
verifyQuota(ctx, t, client, 4, 4)
145+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
127146

128147
// Next one must fail
129148
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
130149
build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
131150

132151
// Consumed shouldn't bump
133-
verifyQuota(ctx, t, client, 4, 4)
152+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
134153
require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
135154
require.Contains(t, build.Job.Error, "quota")
136155

@@ -146,15 +165,15 @@ func TestWorkspaceQuota(t *testing.T) {
146165
})
147166
require.NoError(t, err)
148167
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
149-
verifyQuota(ctx, t, client, 3, 4)
168+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 3, 4)
150169
break
151170
}
152171

153172
// Next one should now succeed
154173
workspace = coderdtest.CreateWorkspace(t, client, template.ID)
155174
build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
156175

157-
verifyQuota(ctx, t, client, 4, 4)
176+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
158177
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
159178
})
160179

@@ -174,14 +193,14 @@ func TestWorkspaceQuota(t *testing.T) {
174193
})
175194
coderdtest.NewProvisionerDaemon(t, api.AGPL)
176195

177-
verifyQuota(ctx, t, client, 0, 0)
196+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0)
178197

179198
// Patch the 'Everyone' group to verify its quota allowance is being accounted for.
180199
_, err := client.PatchGroup(ctx, user.OrganizationID, codersdk.PatchGroupRequest{
181200
QuotaAllowance: ptr.Ref(4),
182201
})
183202
require.NoError(t, err)
184-
verifyQuota(ctx, t, client, 0, 4)
203+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 4)
185204

186205
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
187206
Parse: echo.ParseComplete,
@@ -208,29 +227,29 @@ func TestWorkspaceQuota(t *testing.T) {
208227
assert.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
209228
}
210229
wg.Wait()
211-
verifyQuota(ctx, t, client, 4, 4)
230+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
212231

213232
// Next one must fail
214233
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
215234
build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
216235
require.Contains(t, build.Job.Error, "quota")
217236

218237
// Consumed shouldn't bump
219-
verifyQuota(ctx, t, client, 4, 4)
238+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
220239
require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
221240

222241
build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStop)
223242
build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
224243

225244
// Quota goes down one
226-
verifyQuota(ctx, t, client, 3, 4)
245+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 3, 4)
227246
require.Equal(t, codersdk.WorkspaceStatusStopped, build.Status)
228247

229248
build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStart)
230249
build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
231250

232251
// Quota goes back up
233-
verifyQuota(ctx, t, client, 4, 4)
252+
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
234253
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
235254
})
236255

@@ -273,13 +292,27 @@ func TestWorkspaceQuota(t *testing.T) {
273292
})
274293
require.NoError(t, err)
275294

276-
verifyQuota(ctx, t, member, 0, 30)
277-
// This currently reports the total site wide quotas. We might want to
278-
// org scope this api call in the future.
279-
verifyQuota(ctx, t, owner, 0, 45)
295+
verifyQuota(ctx, t, member, first.OrganizationID.String(), 0, 30)
296+
297+
// Verify org scoped quota limits
298+
verifyQuota(ctx, t, owner, first.OrganizationID.String(), 0, 30)
299+
verifyQuota(ctx, t, owner, second.ID.String(), 0, 15)
280300
})
281301
}
282302

303+
func deprecatedQuotaEndpoint(ctx context.Context, client *codersdk.Client, userID string) (codersdk.WorkspaceQuota, error) {
304+
res, err := client.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil)
305+
if err != nil {
306+
return codersdk.WorkspaceQuota{}, err
307+
}
308+
defer res.Body.Close()
309+
if res.StatusCode != http.StatusOK {
310+
return codersdk.WorkspaceQuota{}, codersdk.ReadBodyAsError(res)
311+
}
312+
var quota codersdk.WorkspaceQuota
313+
return quota, json.NewDecoder(res.Body).Decode(&quota)
314+
}
315+
283316
func planWithCost(cost int32) []*proto.Response {
284317
return []*proto.Response{{
285318
Type: &proto.Response_Plan{

0 commit comments

Comments
 (0)