Chore: Refactor quota service (#57586)

* Chore: refactore quota service

* Apply suggestions from code review

pkg/tests/api/alerting/api_alertmanager_test.go

@@ -16,7 +16,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"github.com/grafana/grafana/pkg/models"
 	apimodels "github.com/grafana/grafana/pkg/services/ngalert/api/tooling/definitions"
 	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
 	ngstore "github.com/grafana/grafana/pkg/services/ngalert/store"
@@ -1878,6 +1877,8 @@ func TestQuota(t *testing.T) {
 
 	// Create a user to make authenticated requests
 	createUser(t, store, user.CreateUserCommand{
+		// needs permission to update org quota
+		IsAdmin:        true,
 		DefaultOrgRole: string(org.RoleEditor),
 		Password:       "password",
 		Login:          "grafana",
@@ -1918,30 +1919,10 @@ func TestQuota(t *testing.T) {
 	// check quota limits
 	t.Run("when quota limit exceed creating new rule should fail", func(t *testing.T) {
 		// get existing org quota
-		query := models.GetOrgQuotaByTargetQuery{OrgId: 1, Target: "alert_rule"}
-		err = store.GetOrgQuotaByTarget(context.Background(), &query)
-		require.NoError(t, err)
-		used := query.Result.Used
-		limit := query.Result.Limit
-
-		// set org quota limit to equal used
-		orgCmd := models.UpdateOrgQuotaCmd{
-			OrgId:  1,
-			Target: "alert_rule",
-			Limit:  used,
-		}
-		err := store.UpdateOrgQuota(context.Background(), &orgCmd)
-		require.NoError(t, err)
-
+		limit, used := apiClient.GetOrgQuotaLimits(t, 1)
+		apiClient.UpdateAlertRuleOrgQuota(t, 1, used)
 		t.Cleanup(func() {
-			// reset org quota to original value
-			orgCmd := models.UpdateOrgQuotaCmd{
-				OrgId:  1,
-				Target: "alert_rule",
-				Limit:  limit,
-			}
-			err := store.UpdateOrgQuota(context.Background(), &orgCmd)
-			require.NoError(t, err)
+			apiClient.UpdateAlertRuleOrgQuota(t, 1, limit)
 		})
 
 		// try to create an alert rule

pkg/tests/api/alerting/testing.go

@@ -16,6 +16,7 @@ import (
 
 	apimodels "github.com/grafana/grafana/pkg/services/ngalert/api/tooling/definitions"
 	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/quota"
 	"github.com/grafana/grafana/pkg/util"
 )
 
@@ -202,6 +203,60 @@ func (a apiClient) CreateFolder(t *testing.T, uID string, title string) {
 	a.ReloadCachedPermissions(t)
 }
 
+func (a apiClient) GetOrgQuotaLimits(t *testing.T, orgID int64) (int64, int64) {
+	t.Helper()
+
+	u := fmt.Sprintf("%s/api/orgs/%d/quotas", a.url, orgID)
+	// nolint:gosec
+	resp, err := http.Get(u)
+	require.NoError(t, err)
+	defer func() {
+		_ = resp.Body.Close()
+	}()
+	b, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	results := []quota.QuotaDTO{}
+	require.NoError(t, json.Unmarshal(b, &results))
+
+	var limit int64 = 0
+	var used int64 = 0
+	for _, q := range results {
+		if q.Target != string(ngmodels.QuotaTargetSrv) {
+			continue
+		}
+		limit = q.Limit
+		used = q.Used
+	}
+	return limit, used
+}
+
+func (a apiClient) UpdateAlertRuleOrgQuota(t *testing.T, orgID int64, limit int64) {
+	t.Helper()
+	buf := bytes.Buffer{}
+	enc := json.NewEncoder(&buf)
+	err := enc.Encode(&quota.UpdateQuotaCmd{
+		Target: "alert_rule",
+		Limit:  limit,
+		OrgID:  orgID,
+	})
+	require.NoError(t, err)
+
+	u := fmt.Sprintf("%s/api/orgs/%d/quotas/alert_rule", a.url, orgID)
+	// nolint:gosec
+	client := &http.Client{}
+	req, err := http.NewRequest(http.MethodPut, u, &buf)
+	require.NoError(t, err)
+	req.Header.Add("Content-Type", "application/json")
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	defer func() {
+		_ = resp.Body.Close()
+	}()
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}
+
 func (a apiClient) PostRulesGroup(t *testing.T, folder string, group *apimodels.PostableRuleGroupConfig) (int, string) {
 	t.Helper()
 	buf := bytes.Buffer{}