mirror of
https://github.com/gitpod-io/gitpod.git
synced 2025-12-08 17:36:30 +00:00
256c896891
when calling `gorm.Save` concurrently, we see duplicate entry errors because gorm is only trying to update and if that doe not succeed does an insert.
469 lines
16 KiB
Go
469 lines
16 KiB
Go
// Copyright (c) 2022 Gitpod GmbH. All rights reserved.
|
|
// Licensed under the GNU Affero General Public License (AGPL).
|
|
// See License.AGPL.txt in the project root for license information.
|
|
|
|
package db_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
db "github.com/gitpod-io/gitpod/components/gitpod-db/go"
|
|
|
|
"github.com/gitpod-io/gitpod/components/gitpod-db/go/dbtest"
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func TestCostCenter_WriteRead(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
|
|
costCenter := &db.CostCenter{
|
|
ID: db.NewTeamAttributionID(uuid.New().String()),
|
|
SpendingLimit: 100,
|
|
}
|
|
cleanUp(t, conn, costCenter.ID)
|
|
|
|
tx := conn.Create(costCenter)
|
|
require.NoError(t, tx.Error)
|
|
|
|
read := &db.CostCenter{ID: costCenter.ID}
|
|
tx = conn.First(read)
|
|
require.NoError(t, tx.Error)
|
|
require.Equal(t, costCenter.ID, read.ID)
|
|
require.Equal(t, costCenter.SpendingLimit, read.SpendingLimit)
|
|
}
|
|
|
|
func TestCostCenterManager_GetOrCreateCostCenter_concurrent(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
id := db.NewTeamAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, id)
|
|
|
|
waitgroup := &sync.WaitGroup{}
|
|
save := func() {
|
|
_, err := mnr.GetOrCreateCostCenter(context.Background(), id)
|
|
require.NoError(t, err)
|
|
waitgroup.Done()
|
|
}
|
|
waitgroup.Add(10)
|
|
for i := 0; i < 10; i++ {
|
|
go save()
|
|
}
|
|
waitgroup.Wait()
|
|
}
|
|
|
|
func TestCostCenterManager_GetOrCreateCostCenter(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
team := db.NewTeamAttributionID(uuid.New().String())
|
|
user := db.NewUserAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, team, user)
|
|
|
|
teamCC, err := mnr.GetOrCreateCostCenter(context.Background(), team)
|
|
require.NoError(t, err)
|
|
userCC, err := mnr.GetOrCreateCostCenter(context.Background(), user)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() {
|
|
conn.Model(&db.CostCenter{}).Delete(teamCC, userCC)
|
|
})
|
|
require.Equal(t, int32(0), teamCC.SpendingLimit)
|
|
require.Equal(t, int32(500), userCC.SpendingLimit)
|
|
}
|
|
|
|
func TestCostCenterManager_GetOrCreateCostCenter_ResetsExpired(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
|
|
now := time.Now().UTC()
|
|
ts := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second(), 0, time.UTC)
|
|
expired := ts.Add(-1 * time.Minute)
|
|
unexpired := ts.Add(1 * time.Minute)
|
|
|
|
expiredCC := db.CostCenter{
|
|
ID: db.NewTeamAttributionID(uuid.New().String()),
|
|
CreationTime: db.NewVarCharTime(now),
|
|
SpendingLimit: 0,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
NextBillingTime: db.NewVarCharTime(expired),
|
|
BillingCycleStart: db.NewVarCharTime(now),
|
|
}
|
|
unexpiredCC := db.CostCenter{
|
|
ID: db.NewUserAttributionID(uuid.New().String()),
|
|
CreationTime: db.NewVarCharTime(now),
|
|
SpendingLimit: 500,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
NextBillingTime: db.NewVarCharTime(unexpired),
|
|
BillingCycleStart: db.NewVarCharTime(now),
|
|
}
|
|
// Stripe billing strategy should not be reset
|
|
stripeCC := db.CostCenter{
|
|
ID: db.NewUserAttributionID(uuid.New().String()),
|
|
CreationTime: db.NewVarCharTime(now),
|
|
SpendingLimit: 0,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
NextBillingTime: db.VarcharTime{},
|
|
BillingCycleStart: db.NewVarCharTime(now),
|
|
}
|
|
|
|
dbtest.CreateCostCenters(t, conn,
|
|
dbtest.NewCostCenter(t, expiredCC),
|
|
dbtest.NewCostCenter(t, unexpiredCC),
|
|
dbtest.NewCostCenter(t, stripeCC),
|
|
)
|
|
|
|
// expired db.CostCenter should be reset, so we get a new CreationTime
|
|
retrievedExpiredCC, err := mnr.GetOrCreateCostCenter(context.Background(), expiredCC.ID)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() {
|
|
conn.Model(&db.CostCenter{}).Delete(retrievedExpiredCC.ID)
|
|
})
|
|
require.Equal(t, db.NewVarCharTime(expired).Time().AddDate(0, 1, 0), retrievedExpiredCC.NextBillingTime.Time())
|
|
require.Equal(t, expiredCC.ID, retrievedExpiredCC.ID)
|
|
require.Equal(t, expiredCC.BillingStrategy, retrievedExpiredCC.BillingStrategy)
|
|
require.WithinDuration(t, now, expiredCC.CreationTime.Time(), 3*time.Second, "new cost center creation time must be within 3 seconds of now")
|
|
|
|
// unexpired cost center must not be reset
|
|
retrievedUnexpiredCC, err := mnr.GetOrCreateCostCenter(context.Background(), unexpiredCC.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, db.NewVarCharTime(unexpired).Time(), retrievedUnexpiredCC.NextBillingTime.Time())
|
|
require.Equal(t, unexpiredCC.ID, retrievedUnexpiredCC.ID)
|
|
require.Equal(t, unexpiredCC.BillingStrategy, retrievedUnexpiredCC.BillingStrategy)
|
|
require.WithinDuration(t, unexpiredCC.CreationTime.Time(), retrievedUnexpiredCC.CreationTime.Time(), 100*time.Millisecond)
|
|
|
|
// stripe cost center must not be reset
|
|
retrievedStripeCC, err := mnr.GetOrCreateCostCenter(context.Background(), stripeCC.ID)
|
|
require.NoError(t, err)
|
|
require.False(t, retrievedStripeCC.NextBillingTime.IsSet())
|
|
require.Equal(t, stripeCC.ID, retrievedStripeCC.ID)
|
|
require.Equal(t, stripeCC.BillingStrategy, retrievedStripeCC.BillingStrategy)
|
|
require.WithinDuration(t, stripeCC.CreationTime.Time(), retrievedStripeCC.CreationTime.Time(), 100*time.Millisecond)
|
|
}
|
|
|
|
func TestCostCenterManager_UpdateCostCenter(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
limits := db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
MinForUsersOnStripe: 1000,
|
|
}
|
|
|
|
t.Run("prevents updates to negative spending limit", func(t *testing.T) {
|
|
mnr := db.NewCostCenterManager(conn, limits)
|
|
userAttributionID := db.NewUserAttributionID(uuid.New().String())
|
|
teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, userAttributionID, teamAttributionID)
|
|
|
|
_, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: userAttributionID,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
SpendingLimit: -1,
|
|
})
|
|
require.Error(t, err)
|
|
require.Equal(t, codes.InvalidArgument, status.Code(err))
|
|
|
|
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: teamAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: -1,
|
|
})
|
|
require.Error(t, err)
|
|
require.Equal(t, codes.InvalidArgument, status.Code(err))
|
|
})
|
|
|
|
t.Run("individual user on Other billing strategy can change spending limit of 500", func(t *testing.T) {
|
|
mnr := db.NewCostCenterManager(conn, limits)
|
|
userAttributionID := db.NewUserAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, userAttributionID)
|
|
|
|
newCC, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: userAttributionID,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
SpendingLimit: 501,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, int32(501), newCC.SpendingLimit)
|
|
|
|
})
|
|
|
|
t.Run("individual user upgrading to stripe can set a limit of 1000 or more, but not less than 1000", func(t *testing.T) {
|
|
mnr := db.NewCostCenterManager(conn, limits)
|
|
userAttributionID := db.NewUserAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, userAttributionID)
|
|
|
|
// Upgrading to Stripe requires spending limit
|
|
res, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: userAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: 1000,
|
|
})
|
|
require.NoError(t, err)
|
|
requireCostCenterEqual(t, db.CostCenter{
|
|
ID: userAttributionID,
|
|
SpendingLimit: 1000,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
}, res)
|
|
|
|
// Try to lower the spending limit below configured limit
|
|
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: userAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: 999,
|
|
})
|
|
require.Error(t, err, "lowering spending limit below configured value is not allowed for user subscriptions")
|
|
|
|
// Try to update the cost center to higher usage limit
|
|
res, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: userAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: 1001,
|
|
})
|
|
require.NoError(t, err)
|
|
requireCostCenterEqual(t, db.CostCenter{
|
|
ID: userAttributionID,
|
|
SpendingLimit: 1001,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
}, res)
|
|
})
|
|
|
|
t.Run("team on Other billing strategy get a spending limit of 0", func(t *testing.T) {
|
|
mnr := db.NewCostCenterManager(conn, limits)
|
|
teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, teamAttributionID)
|
|
|
|
// Allows udpating cost center as long as spending limit remains as configured
|
|
res, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: teamAttributionID,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
SpendingLimit: limits.ForTeams,
|
|
})
|
|
require.NoError(t, err)
|
|
requireCostCenterEqual(t, db.CostCenter{
|
|
ID: teamAttributionID,
|
|
SpendingLimit: limits.ForTeams,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
}, res)
|
|
})
|
|
|
|
t.Run("team on Stripe billing strategy can set arbitrary positive spending limit", func(t *testing.T) {
|
|
mnr := db.NewCostCenterManager(conn, limits)
|
|
teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, teamAttributionID)
|
|
|
|
// Allows udpating cost center as long as spending limit remains as configured
|
|
res, err := mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: teamAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: limits.ForTeams,
|
|
})
|
|
require.NoError(t, err)
|
|
requireCostCenterEqual(t, db.CostCenter{
|
|
ID: teamAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: limits.ForTeams,
|
|
}, res)
|
|
|
|
// Allows updating cost center to any positive value
|
|
_, err = mnr.UpdateCostCenter(context.Background(), db.CostCenter{
|
|
ID: teamAttributionID,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
SpendingLimit: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
})
|
|
}
|
|
|
|
func TestSaveCostCenterMovedToStripe(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 20,
|
|
ForUsers: 500,
|
|
})
|
|
team := db.NewTeamAttributionID(uuid.New().String())
|
|
cleanUp(t, conn, team)
|
|
teamCC, err := mnr.GetOrCreateCostCenter(context.Background(), team)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int32(20), teamCC.SpendingLimit)
|
|
|
|
teamCC.BillingStrategy = db.CostCenter_Stripe
|
|
teamCC.SpendingLimit = 400050
|
|
teamCC, err = mnr.UpdateCostCenter(context.Background(), teamCC)
|
|
require.NoError(t, err)
|
|
require.Equal(t, db.CostCenter_Stripe, teamCC.BillingStrategy)
|
|
require.Equal(t, teamCC.CreationTime.Time().AddDate(0, 1, 0), teamCC.NextBillingTime.Time())
|
|
require.Equal(t, int32(400050), teamCC.SpendingLimit)
|
|
|
|
teamCC.BillingStrategy = db.CostCenter_Other
|
|
teamCC, err = mnr.UpdateCostCenter(context.Background(), teamCC)
|
|
require.NoError(t, err)
|
|
require.Equal(t, teamCC.CreationTime.Time().AddDate(0, 1, 0).Truncate(time.Second), teamCC.NextBillingTime.Time().Truncate(time.Second))
|
|
require.Equal(t, int32(20), teamCC.SpendingLimit)
|
|
}
|
|
|
|
func cleanUp(t *testing.T, conn *gorm.DB, attributionIds ...db.AttributionID) {
|
|
t.Helper()
|
|
t.Cleanup(func() {
|
|
for _, attributionId := range attributionIds {
|
|
conn.Where("id = ?", string(attributionId)).Delete(&db.CostCenter{})
|
|
conn.Where("attributionId = ?", string(attributionId)).Delete(&db.Usage{})
|
|
}
|
|
})
|
|
}
|
|
|
|
func requireCostCenterEqual(t *testing.T, expected, actual db.CostCenter) {
|
|
t.Helper()
|
|
|
|
// ignore timestamps in comparsion
|
|
require.Equal(t, expected.ID, actual.ID)
|
|
require.EqualValues(t, expected.SpendingLimit, actual.SpendingLimit)
|
|
require.Equal(t, expected.BillingStrategy, actual.BillingStrategy)
|
|
}
|
|
|
|
func TestCostCenter_ListLatestCostCentersWithBillingTimeBefore(t *testing.T) {
|
|
|
|
t.Run("no cost centers found when no data exists", func(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
|
|
ts := time.Date(2022, 10, 10, 10, 10, 10, 10, time.UTC)
|
|
|
|
retrieved, err := mnr.ListManagedCostCentersWithBillingTimeBefore(context.Background(), ts.Add(7*24*time.Hour))
|
|
require.NoError(t, err)
|
|
require.Len(t, retrieved, 0)
|
|
})
|
|
|
|
t.Run("returns the most recent cost center (by creation time)", func(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
|
|
attributionID := uuid.New().String()
|
|
firstCreation := time.Date(2022, 10, 10, 10, 10, 10, 10, time.UTC)
|
|
secondCreation := firstCreation.Add(24 * time.Hour)
|
|
|
|
costCenters := []db.CostCenter{
|
|
dbtest.NewCostCenter(t, db.CostCenter{
|
|
ID: db.NewTeamAttributionID(attributionID),
|
|
SpendingLimit: 100,
|
|
CreationTime: db.NewVarCharTime(firstCreation),
|
|
BillingStrategy: db.CostCenter_Other,
|
|
NextBillingTime: db.NewVarCharTime(firstCreation),
|
|
}),
|
|
dbtest.NewCostCenter(t, db.CostCenter{
|
|
ID: db.NewTeamAttributionID(attributionID),
|
|
SpendingLimit: 100,
|
|
CreationTime: db.NewVarCharTime(secondCreation),
|
|
BillingStrategy: db.CostCenter_Other,
|
|
NextBillingTime: db.NewVarCharTime(secondCreation),
|
|
}),
|
|
}
|
|
|
|
dbtest.CreateCostCenters(t, conn, costCenters...)
|
|
|
|
retrieved, err := mnr.ListManagedCostCentersWithBillingTimeBefore(context.Background(), secondCreation.Add(7*24*time.Hour))
|
|
require.NoError(t, err)
|
|
require.Len(t, retrieved, 1)
|
|
|
|
requireCostCenterEqual(t, costCenters[1], retrieved[0])
|
|
})
|
|
|
|
t.Run("returns results only when most recent cost center matches billing strategy", func(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
|
|
attributionID := uuid.New().String()
|
|
firstCreation := time.Date(2022, 10, 10, 10, 10, 10, 10, time.UTC)
|
|
secondCreation := firstCreation.Add(24 * time.Hour)
|
|
|
|
costCenters := []db.CostCenter{
|
|
dbtest.NewCostCenter(t, db.CostCenter{
|
|
ID: db.NewTeamAttributionID(attributionID),
|
|
SpendingLimit: 100,
|
|
CreationTime: db.NewVarCharTime(firstCreation),
|
|
BillingStrategy: db.CostCenter_Other,
|
|
NextBillingTime: db.NewVarCharTime(firstCreation),
|
|
}),
|
|
dbtest.NewCostCenter(t, db.CostCenter{
|
|
ID: db.NewTeamAttributionID(attributionID),
|
|
SpendingLimit: 100,
|
|
CreationTime: db.NewVarCharTime(secondCreation),
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
}),
|
|
}
|
|
|
|
dbtest.CreateCostCenters(t, conn, costCenters...)
|
|
|
|
retrieved, err := mnr.ListManagedCostCentersWithBillingTimeBefore(context.Background(), secondCreation.Add(7*24*time.Hour))
|
|
require.NoError(t, err)
|
|
require.Len(t, retrieved, 0)
|
|
})
|
|
}
|
|
|
|
func TestCostCenterManager_ResetUsage(t *testing.T) {
|
|
now := time.Now().UTC()
|
|
ts := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second(), 0, time.UTC)
|
|
|
|
t.Run("errors when cost center is not Other", func(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
cc := dbtest.CreateCostCenters(t, conn, db.CostCenter{
|
|
ID: db.NewUserAttributionID(uuid.New().String()),
|
|
CreationTime: db.NewVarCharTime(time.Now()),
|
|
SpendingLimit: 500,
|
|
BillingStrategy: db.CostCenter_Stripe,
|
|
})[0]
|
|
|
|
_, err := mnr.ResetUsage(context.Background(), cc.ID)
|
|
require.Error(t, err)
|
|
})
|
|
|
|
t.Run("resets for teams", func(t *testing.T) {
|
|
conn := dbtest.ConnectForTests(t)
|
|
mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
|
|
ForTeams: 0,
|
|
ForUsers: 500,
|
|
})
|
|
oldCC := dbtest.CreateCostCenters(t, conn, db.CostCenter{
|
|
ID: db.NewTeamAttributionID(uuid.New().String()),
|
|
CreationTime: db.NewVarCharTime(time.Now()),
|
|
SpendingLimit: 10,
|
|
BillingStrategy: db.CostCenter_Other,
|
|
NextBillingTime: db.NewVarCharTime(ts),
|
|
BillingCycleStart: db.NewVarCharTime(ts.AddDate(0, -1, 0)),
|
|
})[0]
|
|
newCC, err := mnr.ResetUsage(context.Background(), oldCC.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, oldCC.ID, newCC.ID)
|
|
require.EqualValues(t, 10, newCC.SpendingLimit)
|
|
require.Equal(t, db.CostCenter_Other, newCC.BillingStrategy)
|
|
require.Equal(t, db.NewVarCharTime(ts.AddDate(0, 1, 0)).Time(), newCC.NextBillingTime.Time())
|
|
|
|
})
|
|
|
|
}
|