Skip to content

Commit 47784ab

Browse files
easyCZeasyCZ
committed
[usage] Validate spending limits in UpdateCostCenter
1 parent f810810 commit 47784ab

File tree

2 files changed

+222
-41
lines changed

2 files changed

+222
-41
lines changed

Diff for: ‎components/usage/pkg/db/cost_center.go

+78-25
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import (
1212

1313
"github.com/gitpod-io/gitpod/common-go/log"
1414
"github.com/google/uuid"
15+
"google.golang.org/grpc/codes"
16+
"google.golang.org/grpc/status"
1517
"gorm.io/gorm"
1618
)
1719

@@ -39,8 +41,9 @@ func (d *CostCenter) TableName() string {
3941
}
4042

4143
type DefaultSpendingLimit struct {
42-
ForTeams int32 `json:"forTeams"`
43-
ForUsers int32 `json:"forUsers"`
44+
ForTeams int32 `json:"forTeams"`
45+
ForUsers int32 `json:"forUsers"`
46+
MinForIndividualsOnStripe int32 `json:"minForIndividualsOnStripe`
4447
}
4548

4649
func NewCostCenterManager(conn *gorm.DB, cfg DefaultSpendingLimit) *CostCenterManager {
@@ -102,27 +105,54 @@ func getCostCenter(ctx context.Context, conn *gorm.DB, attributionId Attribution
102105
return costCenter, nil
103106
}
104107

105-
func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, costCenter CostCenter) (CostCenter, error) {
108+
func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, newCC CostCenter) (CostCenter, error) {
109+
if newCC.SpendingLimit < 0 {
110+
return CostCenter{}, status.Errorf(codes.InvalidArgument, "Spending limit cannot be set below zero.")
111+
}
106112

113+
attributionID := newCC.ID
107114
// retrieving the existing cost center to maintain the readonly values
108-
existingCostCenter, err := c.GetOrCreateCostCenter(ctx, costCenter.ID)
115+
existingCC, err := c.GetOrCreateCostCenter(ctx, newCC.ID)
109116
if err != nil {
110-
return CostCenter{}, err
117+
return CostCenter{}, status.Errorf(codes.NotFound, "cost center does not exist")
111118
}
112119

113120
now := time.Now()
114121

115122
// we always update the creationTime
116-
costCenter.CreationTime = NewVarcharTime(now)
123+
newCC.CreationTime = NewVarcharTime(now)
117124
// we don't allow setting the nextBillingTime from outside
118-
costCenter.NextBillingTime = existingCostCenter.NextBillingTime
125+
newCC.NextBillingTime = existingCC.NextBillingTime
126+
127+
isTeam := attributionID.IsEntity(AttributionEntity_Team)
128+
isUser := attributionID.IsEntity(AttributionEntity_User)
129+
130+
if isUser {
131+
// New billing strategy is Stripe
132+
if newCC.BillingStrategy == CostCenter_Stripe {
133+
if newCC.SpendingLimit < c.cfg.MinForIndividualsOnStripe {
134+
return CostCenter{}, status.Errorf(codes.FailedPrecondition, "individual users cannot lower their spending below %d", c.cfg.ForUsers)
135+
}
136+
}
137+
138+
// Billing strategy remains unchanged
139+
if newCC.BillingStrategy == CostCenter_Other && existingCC.BillingStrategy == CostCenter_Other {
140+
if newCC.SpendingLimit != existingCC.SpendingLimit {
141+
return CostCenter{}, status.Errorf(codes.FailedPrecondition, "individual users on a free plan cannot adjust their spending limit")
142+
}
143+
}
119144

120-
// Do we have a billing strategy update?
121-
if costCenter.BillingStrategy != existingCostCenter.BillingStrategy {
122-
switch costCenter.BillingStrategy {
123-
case CostCenter_Stripe:
145+
// Downgrading from stripe
146+
if existingCC.BillingStrategy == CostCenter_Stripe && newCC.BillingStrategy == CostCenter_Other {
147+
newCC.SpendingLimit = c.cfg.ForUsers
148+
// see you next month
149+
newCC.NextBillingTime = NewVarcharTime(now.AddDate(0, 1, 0))
150+
}
151+
152+
// Upgrading to Stripe
153+
if existingCC.BillingStrategy == CostCenter_Other && newCC.BillingStrategy == CostCenter_Stripe {
124154
// moving to stripe -> let's run a finalization
125-
finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, costCenter.ID)
155+
finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, newCC.ID)
126156
if err != nil {
127157
return CostCenter{}, err
128158
}
@@ -133,26 +163,49 @@ func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, costCenter Cos
133163
}
134164
}
135165
// we don't manage stripe billing cycle
136-
costCenter.NextBillingTime = VarcharTime{}
137-
138-
case CostCenter_Other:
139-
// cancelled from stripe reset the spending limit
140-
if costCenter.ID.IsEntity(AttributionEntity_Team) {
141-
costCenter.SpendingLimit = c.cfg.ForTeams
142-
} else {
143-
costCenter.SpendingLimit = c.cfg.ForUsers
166+
newCC.NextBillingTime = VarcharTime{}
167+
}
168+
} else if isTeam {
169+
// Billing strategy is Other, and it remains unchanged
170+
if existingCC.BillingStrategy == CostCenter_Other && newCC.BillingStrategy == CostCenter_Other {
171+
// It is impossible for a team without Stripe billing to change their spending limit
172+
if newCC.SpendingLimit != c.cfg.ForTeams {
173+
return CostCenter{}, status.Errorf(codes.FailedPrecondition, "teams without a subscription cannot change their spending limit")
144174
}
175+
}
176+
177+
// Downgrading from stripe
178+
if existingCC.BillingStrategy == CostCenter_Stripe && newCC.BillingStrategy == CostCenter_Other {
179+
newCC.SpendingLimit = c.cfg.ForTeams
145180
// see you next month
146-
costCenter.NextBillingTime = NewVarcharTime(now.AddDate(0, 1, 0))
181+
newCC.NextBillingTime = NewVarcharTime(now.AddDate(0, 1, 0))
147182
}
183+
184+
// Upgrading to Stripe
185+
if existingCC.BillingStrategy == CostCenter_Other && newCC.BillingStrategy == CostCenter_Stripe {
186+
// moving to stripe -> let's run a finalization
187+
finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, newCC.ID)
188+
if err != nil {
189+
return CostCenter{}, err
190+
}
191+
if finalizationUsage != nil {
192+
err = UpdateUsage(ctx, c.conn, *finalizationUsage)
193+
if err != nil {
194+
return CostCenter{}, err
195+
}
196+
}
197+
// we don't manage stripe billing cycle
198+
newCC.NextBillingTime = VarcharTime{}
199+
}
200+
148201
}
149202

150-
log.WithField("cost_center", costCenter).Info("saving cost center.")
151-
db := c.conn.Save(&costCenter)
203+
log.WithField("cost_center", newCC).Info("saving cost center.")
204+
db := c.conn.Save(&newCC)
152205
if db.Error != nil {
153-
return CostCenter{}, fmt.Errorf("failed to save cost center for attributionID %s: %w", costCenter.ID, db.Error)
206+
return CostCenter{}, fmt.Errorf("failed to save cost center for attributionID %s: %w", newCC.ID, db.Error)
154207
}
155-
return costCenter, nil
208+
return newCC, nil
156209
}
157210

158211
func (c *CostCenterManager) ComputeInvoiceUsageRecord(ctx context.Context, attributionID AttributionID) (*Usage, error) {

Diff for: ‎components/usage/pkg/db/cost_center_test.go

+144-16
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import (
1313
"github.com/gitpod-io/gitpod/usage/pkg/db/dbtest"
1414
"github.com/google/uuid"
1515
"github.com/stretchr/testify/require"
16+
"google.golang.org/grpc/codes"
17+
"google.golang.org/grpc/status"
1618
"gorm.io/gorm"
1719
)
1820

@@ -58,26 +60,143 @@ func TestCostCenterManager_GetOrCreateCostCenter(t *testing.T) {
5860

5961
func TestCostCenterManager_UpdateCostCenter(t *testing.T) {
6062
conn := dbtest.ConnectForTests(t)
61-
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
62-
ForTeams: 0,
63-
ForUsers: 500,
63+
limits := db.DefaultSpendingLimit{
64+
ForTeams: 0,
65+
ForUsers: 500,
66+
MinForIndividualsOnStripe: 1000,
67+
}
68+
69+
t.Run("prevents updates to negative spending limit", func(t *testing.T) {
70+
mnr := db.NewCostCenterManager(conn, limits)
71+
userAttributionID := db.NewUserAttributionID(uuid.New().String())
72+
teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
73+
cleanUp(t, conn, userAttributionID, teamAttributionID)
74+
75+
_, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
76+
ID: userAttributionID,
77+
BillingStrategy: db.CostCenter_Other,
78+
SpendingLimit: -1,
79+
})
80+
require.Error(t, err)
81+
require.Equal(t, codes.InvalidArgument, status.Code(err))
82+
83+
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
84+
ID: teamAttributionID,
85+
BillingStrategy: db.CostCenter_Stripe,
86+
SpendingLimit: -1,
87+
})
88+
require.Error(t, err)
89+
require.Equal(t, codes.InvalidArgument, status.Code(err))
6490
})
65-
team := db.NewTeamAttributionID(uuid.New().String())
66-
cleanUp(t, conn, team)
67-
teamCC, err := mnr.GetOrCreateCostCenter(context.Background(), team)
68-
t.Cleanup(func() {
69-
conn.Model(&db.CostCenter{}).Delete(teamCC)
91+
92+
t.Run("individual user on Other billing strategy cannot change spending limit of 500", func(t *testing.T) {
93+
mnr := db.NewCostCenterManager(conn, limits)
94+
userAttributionID := db.NewUserAttributionID(uuid.New().String())
95+
cleanUp(t, conn, userAttributionID)
96+
97+
_, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
98+
ID: userAttributionID,
99+
BillingStrategy: db.CostCenter_Other,
100+
SpendingLimit: 501,
101+
})
102+
require.Error(t, err)
103+
require.Equal(t, codes.FailedPrecondition, status.Code(err))
104+
70105
})
71-
require.NoError(t, err)
72-
require.Equal(t, int32(0), teamCC.SpendingLimit)
73106

74-
teamCC.SpendingLimit = 2000
75-
teamCC, err = mnr.UpdateCostCenter(context.Background(), teamCC)
76-
require.NoError(t, err)
77-
t.Cleanup(func() {
78-
conn.Model(&db.CostCenter{}).Delete(teamCC)
107+
t.Run("individual user upgrading to stripe gets default limit of 1000 and can increase, but not decrease", func(t *testing.T) {
108+
mnr := db.NewCostCenterManager(conn, limits)
109+
userAttributionID := db.NewUserAttributionID(uuid.New().String())
110+
cleanUp(t, conn, userAttributionID)
111+
112+
// Upgrading to Stripe requires spending limit
113+
res, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
114+
ID: userAttributionID,
115+
BillingStrategy: db.CostCenter_Stripe,
116+
SpendingLimit: 1000,
117+
})
118+
require.NoError(t, err)
119+
requireCostCenterEqual(t, db.CostCenter{
120+
ID: userAttributionID,
121+
SpendingLimit: 1000,
122+
BillingStrategy: db.CostCenter_Stripe,
123+
}, res)
124+
125+
// Try to lower the spending limit below configured limit
126+
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
127+
ID: userAttributionID,
128+
BillingStrategy: db.CostCenter_Stripe,
129+
SpendingLimit: 999,
130+
})
131+
require.Error(t, err, "lowering spending limit below configured value is not allowed for user subscriptions")
132+
133+
// Try to update the cost center to higher usage limit
134+
res, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
135+
ID: userAttributionID,
136+
BillingStrategy: db.CostCenter_Stripe,
137+
SpendingLimit: 1001,
138+
})
139+
require.NoError(t, err)
140+
requireCostCenterEqual(t, db.CostCenter{
141+
ID: userAttributionID,
142+
SpendingLimit: 1001,
143+
BillingStrategy: db.CostCenter_Stripe,
144+
}, res)
145+
})
146+
147+
t.Run("team on Other billing strategy get a spending limit of 0, and cannot change it", func(t *testing.T) {
148+
mnr := db.NewCostCenterManager(conn, limits)
149+
teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
150+
cleanUp(t, conn, teamAttributionID)
151+
152+
// Allows udpating cost center as long as spending limit remains as configured
153+
res, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
154+
ID: teamAttributionID,
155+
BillingStrategy: db.CostCenter_Other,
156+
SpendingLimit: limits.ForTeams,
157+
})
158+
require.NoError(t, err)
159+
requireCostCenterEqual(t, db.CostCenter{
160+
ID: teamAttributionID,
161+
SpendingLimit: limits.ForTeams,
162+
BillingStrategy: db.CostCenter_Other,
163+
}, res)
164+
165+
// Prevents updating when spending limit changes
166+
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
167+
ID: teamAttributionID,
168+
BillingStrategy: db.CostCenter_Other,
169+
SpendingLimit: 1,
170+
})
171+
require.Error(t, err)
172+
})
173+
174+
t.Run("team on Stripe billing strategy can set arbitrary positive spending limit", func(t *testing.T) {
175+
mnr := db.NewCostCenterManager(conn, limits)
176+
teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
177+
cleanUp(t, conn, teamAttributionID)
178+
179+
// Allows udpating cost center as long as spending limit remains as configured
180+
res, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
181+
ID: teamAttributionID,
182+
BillingStrategy: db.CostCenter_Stripe,
183+
SpendingLimit: limits.ForTeams,
184+
})
185+
require.NoError(t, err)
186+
requireCostCenterEqual(t, db.CostCenter{
187+
ID: teamAttributionID,
188+
BillingStrategy: db.CostCenter_Stripe,
189+
SpendingLimit: limits.ForTeams,
190+
}, res)
191+
192+
// Allows updating cost center to any positive value
193+
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
194+
ID: teamAttributionID,
195+
BillingStrategy: db.CostCenter_Stripe,
196+
SpendingLimit: 10,
197+
})
198+
require.NoError(t, err)
79199
})
80-
require.Equal(t, int32(2000), teamCC.SpendingLimit)
81200
}
82201

83202
func TestSaveCostCenterMovedToStripe(t *testing.T) {
@@ -116,3 +235,12 @@ func cleanUp(t *testing.T, conn *gorm.DB, attributionIds ...db.AttributionID) {
116235
}
117236
})
118237
}
238+
239+
func requireCostCenterEqual(t *testing.T, expected, actual db.CostCenter) {
240+
t.Helper()
241+
242+
// ignore timestamps in comparsion
243+
require.Equal(t, expected.ID, actual.ID)
244+
require.EqualValues(t, expected.SpendingLimit, actual.SpendingLimit)
245+
require.Equal(t, expected.BillingStrategy, actual.BillingStrategy)
246+
}

0 commit comments

Comments
 (0)