Skip to content

Commit f253dc9

Browse files
easyCZeasyCZ
committed
[usage] Validate spending limits in UpdateCostCenter
1 parent c67e609 commit f253dc9

File tree

3 files changed

+244
-57
lines changed

3 files changed

+244
-57
lines changed

Diff for: ‎components/usage/pkg/apiv1/usage_test.go

+9-8
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,9 @@ func newUsageService(t *testing.T, dbconn *gorm.DB) v1.UsageServiceClient {
8080
)
8181

8282
costCenterManager := db.NewCostCenterManager(dbconn, db.DefaultSpendingLimit{
83-
ForTeams: 0,
84-
ForUsers: 500,
83+
ForTeams: 0,
84+
ForUsers: 500,
85+
MinForIndividualsOnStripe: 1000,
8586
})
8687

8788
v1.RegisterUsageServiceServer(srv.GRPC(), NewUsageService(dbconn, DefaultWorkspacePricer, costCenterManager))
@@ -227,13 +228,13 @@ func TestGetAndSetCostCenter(t *testing.T) {
227228
conn := dbtest.ConnectForTests(t)
228229
costCenterUpdates := []*v1.CostCenter{
229230
{
230-
AttributionId: string(db.NewTeamAttributionID(uuid.New().String())),
231-
SpendingLimit: 5000,
232-
BillingStrategy: v1.CostCenter_BILLING_STRATEGY_OTHER,
231+
AttributionId: string(db.NewUserAttributionID(uuid.New().String())),
232+
SpendingLimit: 8000,
233+
BillingStrategy: v1.CostCenter_BILLING_STRATEGY_STRIPE,
233234
},
234235
{
235-
AttributionId: string(db.NewTeamAttributionID(uuid.New().String())),
236-
SpendingLimit: 8000,
236+
AttributionId: string(db.NewUserAttributionID(uuid.New().String())),
237+
SpendingLimit: 500,
237238
BillingStrategy: v1.CostCenter_BILLING_STRATEGY_OTHER,
238239
},
239240
{
@@ -243,7 +244,7 @@ func TestGetAndSetCostCenter(t *testing.T) {
243244
},
244245
{
245246
AttributionId: string(db.NewTeamAttributionID(uuid.New().String())),
246-
SpendingLimit: 5000,
247+
SpendingLimit: 0,
247248
BillingStrategy: v1.CostCenter_BILLING_STRATEGY_OTHER,
248249
},
249250
}

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

+91-33
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,57 +105,112 @@ 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
119-
120-
// Do we have a billing strategy update?
121-
if costCenter.BillingStrategy != existingCostCenter.BillingStrategy {
122-
switch costCenter.BillingStrategy {
123-
case CostCenter_Stripe:
124-
// moving to stripe -> let's run a finalization
125-
finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, costCenter.ID)
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+
}
144+
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 {
154+
err := c.upgradeToStripe(ctx, attributionID)
126155
if err != nil {
127156
return CostCenter{}, err
128157
}
129-
if finalizationUsage != nil {
130-
err = UpdateUsage(ctx, c.conn, *finalizationUsage)
131-
if err != nil {
132-
return CostCenter{}, err
133-
}
134-
}
158+
135159
// 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
160+
newCC.NextBillingTime = VarcharTime{}
161+
}
162+
} else if isTeam {
163+
// Billing strategy is Other, and it remains unchanged
164+
if existingCC.BillingStrategy == CostCenter_Other && newCC.BillingStrategy == CostCenter_Other {
165+
// It is impossible for a team without Stripe billing to change their spending limit
166+
if newCC.SpendingLimit != c.cfg.ForTeams {
167+
return CostCenter{}, status.Errorf(codes.FailedPrecondition, "teams without a subscription cannot change their spending limit")
144168
}
169+
}
170+
171+
// Downgrading from stripe
172+
if existingCC.BillingStrategy == CostCenter_Stripe && newCC.BillingStrategy == CostCenter_Other {
173+
newCC.SpendingLimit = c.cfg.ForTeams
145174
// see you next month
146-
costCenter.NextBillingTime = NewVarcharTime(now.AddDate(0, 1, 0))
175+
newCC.NextBillingTime = NewVarcharTime(now.AddDate(0, 1, 0))
176+
}
177+
178+
// Upgrading to Stripe
179+
if existingCC.BillingStrategy == CostCenter_Other && newCC.BillingStrategy == CostCenter_Stripe {
180+
err := c.upgradeToStripe(ctx, attributionID)
181+
if err != nil {
182+
return CostCenter{}, err
183+
}
184+
185+
// we don't manage stripe billing cycle
186+
newCC.NextBillingTime = VarcharTime{}
147187
}
188+
} else {
189+
return CostCenter{}, status.Errorf(codes.InvalidArgument, "Unknown attribution entity %s", string(attributionID))
148190
}
149191

150-
log.WithField("cost_center", costCenter).Info("saving cost center.")
151-
db := c.conn.Save(&costCenter)
192+
log.WithField("cost_center", newCC).Info("saving cost center.")
193+
db := c.conn.Save(&newCC)
152194
if db.Error != nil {
153-
return CostCenter{}, fmt.Errorf("failed to save cost center for attributionID %s: %w", costCenter.ID, db.Error)
195+
return CostCenter{}, fmt.Errorf("failed to save cost center for attributionID %s: %w", newCC.ID, db.Error)
154196
}
155-
return costCenter, nil
197+
return newCC, nil
198+
}
199+
200+
func (c *CostCenterManager) upgradeToStripe(ctx context.Context, attributionID AttributionID) error {
201+
// moving to stripe -> let's run a finalization
202+
finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, attributionID)
203+
if err != nil {
204+
return err
205+
}
206+
if finalizationUsage != nil {
207+
err = UpdateUsage(ctx, c.conn, *finalizationUsage)
208+
if err != nil {
209+
return err
210+
}
211+
}
212+
213+
return nil
156214
}
157215

158216
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)