add tests for flags

f0ssel · f0ssel · commit 31c1baa895c4 · 2022-09-29T18:12:00.000Z
diff --git a/codersdk/workspacequota.go b/codersdk/workspacequota.go
@@ -1,6 +1,26 @@
 package codersdk
 
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+)
+
 type WorkspaceQuota struct {
 	UserWorkspaceCount int `json:"user_workspace_count"`
 	UserWorkspaceLimit int `json:"user_workspace_limit"`
 }
+
+func (c *Client) WorkspaceQuota(ctx context.Context, userID string) (WorkspaceQuota, error) {
+	res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/users/%s/workspace-quota", userID), nil)
+	if err != nil {
+		return WorkspaceQuota{}, err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusOK {
+		return WorkspaceQuota{}, readBodyAsError(res)
+	}
+	var quota WorkspaceQuota
+	return quota, json.NewDecoder(res.Body).Decode(&quota)
+}
diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go
@@ -0,0 +1,57 @@
+package coderd_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/enterprise/coderd/coderdenttest"
+	"github.com/coder/coder/testutil"
+	"github.com/stretchr/testify/require"
+)
+
+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, "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, "me")
+		require.NoError(t, err)
+		require.EqualValues(t, q1.UserWorkspaceLimit, max)
+
+		// ensure other user IDs work too
+		u2, err := client.CreateUser(ctx, codersdk.CreateUserRequest{
+			Email:          "whatever@yo.com",
+			Username:       "haha",
+			Password:       "laskjdnvkaj",
+			OrganizationID: user.OrganizationID,
+		})
+		q2, err := client.WorkspaceQuota(ctx, u2.ID.String())
+		require.NoError(t, err)
+		require.EqualValues(t, q1, q2)
+	})
+}
