synced/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2025-07-15 22:20:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement Quotas v3 (#5012)

Browse Source 
* provisioner/terraform: add cost to resource_metadata

* provisionerd/runner: use Options struct

* Complete provisionerd implementation

* Add quota_allowance to groups

* Combine Quota and RBAC licenses

* Add Opts to InTx
This commit is contained in:
Ammar Bandukwala
2022-11-14 11:57:33 -06:00
committed by GitHub
parent 3fb7892c07
commit 97dbd4dc5d
101 changed files with 1577 additions and 847 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
enterprise/coderd/workspacequota_test.go
View File

@ -3,13 +3,11 @@ package coderd_test
import (
"context"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/coder/coder/coderd/coderdtest"
"github.com/coder/coder/coderd/util/ptr"
"github.com/coder/coder/codersdk"
"github.com/coder/coder/enterprise/coderd/coderdenttest"
"github.com/coder/coder/provisioner/echo"
@ -17,49 +15,22 @@ import (
"github.com/coder/coder/testutil"
)
func TestWorkspaceQuota(t *testing.T) {
t.Parallel()
t.Run("Disabled", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
client := coderdenttest.New(t, &coderdenttest.Options{})
_ = coderdtest.CreateFirstUser(t, client)
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
WorkspaceQuota: true,
})
q1, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceLimit, 0)
})
t.Run("Enabled", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
max := 3
client := coderdenttest.New(t, &coderdenttest.Options{
UserWorkspaceQuota: max,
})
user := coderdtest.CreateFirstUser(t, client)
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
WorkspaceQuota: true,
})
q1, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceLimit, max)
func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, consumed, total int) {
t.Helper()
got, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, codersdk.WorkspaceQuota{
Budget: total,
CreditsConsumed: consumed,
}, got)
}
func TestWorkspaceQuota(t *testing.T) {
// TODO: refactor for new impl
t.Parallel()
// ensure other user IDs work too
u2, err := client.CreateUser(ctx, codersdk.CreateUserRequest{
Email: "whatever@yo.com",
Username: "haha",
Password: "laskjdnvkaj",
OrganizationID: user.OrganizationID,
})
require.NoError(t, err)
q2, err := client.WorkspaceQuota(ctx, u2.ID.String())
require.NoError(t, err)
require.EqualValues(t, q1, q2)
})
t.Run("BlocksBuild", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@ -71,14 +42,38 @@ func TestWorkspaceQuota(t *testing.T) {
IncludeProvisionerDaemon: true,
},
})
user := coderdtest.CreateFirstUser(t, client)
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
WorkspaceQuota: true,
TemplateRBAC: true,
})
verifyQuota(ctx, t, client, 0, 0)
// Add user to two groups, granting them a total budget of 3.
group1, err := client.CreateGroup(ctx, user.OrganizationID, codersdk.CreateGroupRequest{
Name: "test-1",
QuotaAllowance: 1,
})
q1, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceCount, 0)
require.EqualValues(t, q1.UserWorkspaceLimit, max)
group2, err := client.CreateGroup(ctx, user.OrganizationID, codersdk.CreateGroupRequest{
Name: "test-2",
QuotaAllowance: 2,
})
require.NoError(t, err)
_, err = client.PatchGroup(ctx, group1.ID, codersdk.PatchGroupRequest{
AddUsers: []string{user.UserID.String()},
})
require.NoError(t, err)
_, err = client.PatchGroup(ctx, group2.ID, codersdk.PatchGroupRequest{
AddUsers: []string{user.UserID.String()},
})
require.NoError(t, err)
verifyQuota(ctx, t, client, 0, 3)
authToken := uuid.NewString()
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
@ -87,8 +82,9 @@ func TestWorkspaceQuota(t *testing.T) {
Type: &proto.Provision_Response_Complete{
Complete: &proto.Provision_Complete{
Resources: []*proto.Resource{{
Name: "example",
Type: "aws_instance",
Name: "example",
Type: "aws_instance",
DailyCost: 1,
Agents: []*proto.Agent{{
Id: uuid.NewString(),
Name: "example",
@ -103,20 +99,45 @@ func TestWorkspaceQuota(t *testing.T) {
})
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
_ = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
_, err = client.CreateWorkspace(context.Background(), user.OrganizationID, codersdk.Me, codersdk.CreateWorkspaceRequest{
TemplateID: template.ID,
Name: "ajksdnvksjd",
AutostartSchedule: ptr.Ref("CRON_TZ=US/Central 30 9 * * 1-5"),
TTLMillis: ptr.Ref((8 * time.Hour).Milliseconds()),
})
require.Error(t, err)
require.ErrorContains(t, err, "User workspace limit")
// ensure count increments
q1, err = client.WorkspaceQuota(ctx, codersdk.Me)
// Spin up three workspaces fine
for i := 0; i < 3; i++ {
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
verifyQuota(ctx, t, client, i+1, 3)
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
}
// Next one must fail
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
// Consumed shouldn't bump
verifyQuota(ctx, t, client, 3, 3)
require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
require.Contains(t, build.Job.Error, "quota")
// Delete one random workspace, then quota should recover.
workspaces, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{})
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceCount, 1)
require.EqualValues(t, q1.UserWorkspaceLimit, max)
for _, w := range workspaces.Workspaces {
if w.LatestBuild.Status != codersdk.WorkspaceStatusRunning {
continue
}
build, err := client.CreateWorkspaceBuild(ctx, w.ID, codersdk.CreateWorkspaceBuildRequest{
Transition: codersdk.WorkspaceTransitionDelete,
})
require.NoError(t, err)
coderdtest.AwaitWorkspaceBuildJob(t, client, build.ID)
verifyQuota(ctx, t, client, 2, 3)
break
}
// Next one should now succeed
workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
build = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
verifyQuota(ctx, t, client, 3, 3)
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
})
}