Chore: Refactor quota service (#57586)
* Chore: refactore quota service * Apply suggestions from code review
This commit is contained in:
Sofia Papagiannaki
@@ -14,15 +14,15 @@ func Quota(quotaService quota.Service) func(string) web.Handler {
|
||||
panic("quotaService is nil")
|
||||
}
|
||||
//https://open.spotify.com/track/7bZSoBEAEEUsGEuLOf94Jm?si=T1Tdju5qRSmmR0zph_6RBw fuuuuunky
|
||||
return func(target string) web.Handler {
|
||||
return func(targetSrv string) web.Handler {
|
||||
return func(c *models.ReqContext) {
|
||||
limitReached, err := quotaService.QuotaReached(c, target)
|
||||
limitReached, err := quotaService.QuotaReached(c, quota.TargetSrv(targetSrv))
|
||||
if err != nil {
|
||||
c.JsonApiErr(500, "Failed to get quota", err)
|
||||
return
|
||||
}
|
||||
if limitReached {
|
||||
c.JsonApiErr(403, fmt.Sprintf("%s Quota reached", target), nil)
|
||||
c.JsonApiErr(403, fmt.Sprintf("%s Quota reached", targetSrv), nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/grafana/grafana/pkg/models"
|
||||
"github.com/grafana/grafana/pkg/services/quota"
|
||||
"github.com/grafana/grafana/pkg/services/quota/quotatest"
|
||||
"github.com/grafana/grafana/pkg/services/user"
|
||||
"github.com/grafana/grafana/pkg/setting"
|
||||
"github.com/grafana/grafana/pkg/web"
|
||||
@@ -30,8 +30,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 403, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Global.User = 4
|
||||
})
|
||||
|
||||
middlewareScenario(t, "and global session quota not reached", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -41,8 +39,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 200, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Global.Session = 10
|
||||
})
|
||||
|
||||
middlewareScenario(t, "and global session quota reached", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -52,13 +48,10 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 403, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Global.Session = 1
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("with user logged in", func(t *testing.T) {
|
||||
const quotaUsed = 4
|
||||
setUp := func(sc *scenarioContext) {
|
||||
sc.withTokenSessionCookie("token")
|
||||
sc.userService.ExpectedSignedInUser = &user.SignedInUser{UserID: 12}
|
||||
@@ -79,8 +72,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 403, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Global.DataSource = quotaUsed
|
||||
})
|
||||
|
||||
middlewareScenario(t, "user Org quota not reached", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -93,8 +84,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 200, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.User.Org = quotaUsed + 1
|
||||
})
|
||||
|
||||
middlewareScenario(t, "user Org quota reached", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -106,8 +95,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 403, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.User.Org = quotaUsed
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org dashboard quota not reached", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -119,8 +106,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 200, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Org.Dashboard = quotaUsed + 1
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org dashboard quota reached", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -132,8 +117,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 403, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Org.Dashboard = quotaUsed
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org dashboard quota reached, but quotas disabled", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -145,9 +128,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 200, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Org.Dashboard = quotaUsed
|
||||
cfg.Quota.Enabled = false
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org alert quota reached and unified alerting is enabled", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -162,7 +142,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
|
||||
cfg.UnifiedAlerting.Enabled = new(bool)
|
||||
*cfg.UnifiedAlerting.Enabled = true
|
||||
cfg.Quota.Org.AlertRule = quotaUsed
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org alert quota not reached and unified alerting is enabled", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -177,7 +156,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
|
||||
cfg.UnifiedAlerting.Enabled = new(bool)
|
||||
*cfg.UnifiedAlerting.Enabled = true
|
||||
cfg.Quota.Org.AlertRule = quotaUsed + 1
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org alert quota reached but ngalert disabled", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -190,8 +168,6 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 403, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Org.AlertRule = quotaUsed
|
||||
})
|
||||
|
||||
middlewareScenario(t, "org alert quota not reached but ngalert disabled", func(t *testing.T, sc *scenarioContext) {
|
||||
@@ -203,58 +179,15 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
assert.Equal(t, 200, sc.resp.Code)
|
||||
}, func(cfg *setting.Cfg) {
|
||||
configure(cfg)
|
||||
|
||||
cfg.Quota.Org.AlertRule = quotaUsed + 1
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func getQuotaHandler(reached bool, target string) web.Handler {
|
||||
qs := &mockQuotaService{
|
||||
reached: reached,
|
||||
}
|
||||
qs := quotatest.New(reached, nil)
|
||||
return Quota(qs)(target)
|
||||
}
|
||||
|
||||
func configure(cfg *setting.Cfg) {
|
||||
cfg.AnonymousEnabled = false
|
||||
cfg.Quota = setting.QuotaSettings{
|
||||
Enabled: true,
|
||||
Org: &setting.OrgQuota{
|
||||
User: 5,
|
||||
Dashboard: 5,
|
||||
DataSource: 5,
|
||||
ApiKey: 5,
|
||||
AlertRule: 5,
|
||||
},
|
||||
User: &setting.UserQuota{
|
||||
Org: 5,
|
||||
},
|
||||
Global: &setting.GlobalQuota{
|
||||
Org: 5,
|
||||
User: 5,
|
||||
Dashboard: 5,
|
||||
DataSource: 5,
|
||||
ApiKey: 5,
|
||||
Session: 5,
|
||||
AlertRule: 5,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type mockQuotaService struct {
|
||||
reached bool
|
||||
err error
|
||||
}
|
||||
|
||||
func (m *mockQuotaService) QuotaReached(c *models.ReqContext, target string) (bool, error) {
|
||||
return m.reached, m.err
|
||||
}
|
||||
|
||||
func (m *mockQuotaService) CheckQuotaReached(c context.Context, target string, params *quota.ScopeParameters) (bool, error) {
|
||||
return m.reached, m.err
|
||||
}
|
||||
|
||||
func (m *mockQuotaService) DeleteByUser(c context.Context, userID int64) error {
|
||||
return m.err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user