// 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 apiv1 import ( "context" "database/sql" "fmt" "testing" "time" "github.com/gitpod-io/gitpod/common-go/baseserver" db "github.com/gitpod-io/gitpod/components/gitpod-db/go" "github.com/gitpod-io/gitpod/components/gitpod-db/go/dbtest" v1 "github.com/gitpod-io/gitpod/usage-api/v1" "github.com/google/uuid" "github.com/stretchr/testify/require" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" "google.golang.org/protobuf/types/known/timestamppb" "gorm.io/gorm" ) func TestUsageService_ReconcileUsage(t *testing.T) { dbconn := dbtest.ConnectForTests(t) from := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC) to := time.Date(2022, 05, 1, 1, 00, 00, 00, time.UTC) attributionID := db.NewTeamAttributionID(uuid.New().String()) t.Cleanup(func() { require.NoError(t, dbconn.Where("attributionId = ?", attributionID).Delete(&db.Usage{}).Error) }) // stopped instances instance := dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ UsageAttributionID: attributionID, StartedTime: db.NewVarCharTime(from), StoppingTime: db.NewVarCharTime(to.Add(-1 * time.Minute)), }) dbtest.CreateWorkspaceInstances(t, dbconn, instance) // running instances dbtest.CreateWorkspaceInstances(t, dbconn, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ StartedTime: db.NewVarCharTime(to.Add(-1 * time.Minute)), UsageAttributionID: attributionID, })) // usage drafts dbtest.CreateUsageRecords(t, dbconn, dbtest.NewUsage(t, db.Usage{ ID: uuid.New(), AttributionID: attributionID, WorkspaceInstanceID: &instance.ID, Kind: db.WorkspaceInstanceUsageKind, Draft: true, })) client := newUsageService(t, dbconn) _, err := client.ReconcileUsage(context.Background(), &v1.ReconcileUsageRequest{ From: timestamppb.New(from), To: timestamppb.New(to), }) require.NoError(t, err) usage, err := db.FindUsage(context.Background(), dbconn, &db.FindUsageParams{ AttributionId: attributionID, From: from, To: to, ExcludeDrafts: false, }) require.NoError(t, err) require.Len(t, usage, 1) } func newUsageService(t *testing.T, dbconn *gorm.DB) v1.UsageServiceClient { srv := baseserver.NewForTests(t, baseserver.WithGRPC(baseserver.MustUseRandomLocalAddress(t)), ) costCenterManager := db.NewCostCenterManager(dbconn, db.DefaultSpendingLimit{ ForTeams: 0, ForUsers: 500, MinForUsersOnStripe: 1000, }) v1.RegisterUsageServiceServer(srv.GRPC(), NewUsageService(dbconn, DefaultWorkspacePricer, costCenterManager)) baseserver.StartServerForTests(t, srv) conn, err := grpc.Dial(srv.GRPCAddress(), grpc.WithTransportCredentials(insecure.NewCredentials())) require.NoError(t, err) client := v1.NewUsageServiceClient(conn) return client } func TestReconcile(t *testing.T) { now := time.Date(2022, 9, 1, 10, 0, 0, 0, time.UTC) pricer, err := NewWorkspacePricer(map[string]float64{ "default": 0.1666666667, "g1-standard": 0.1666666667, "g1-standard-pvc": 0.1666666667, "g1-large": 0.3333333333, "g1-large-pvc": 0.3333333333, "gitpodio-internal-xl": 0.3333333333, }) require.NoError(t, err) t.Run("no action with no instances and no drafts", func(t *testing.T) { inserts, updates, err := reconcileUsage(nil, nil, pricer, now) require.NoError(t, err) require.Len(t, inserts, 0) require.Len(t, updates, 0) }) t.Run("no action with no instances but existing drafts", func(t *testing.T) { drafts := []db.Usage{dbtest.NewUsage(t, db.Usage{})} inserts, updates, err := reconcileUsage(nil, drafts, pricer, now) require.NoError(t, err) require.Len(t, inserts, 0) require.Len(t, updates, 0) }) t.Run("creates a new usage record when no draft exists, removing duplicates", func(t *testing.T) { instance := db.WorkspaceInstanceForUsage{ ID: uuid.New(), WorkspaceID: dbtest.GenerateWorkspaceID(), OwnerID: uuid.New(), ProjectID: sql.NullString{ String: "my-project", Valid: true, }, WorkspaceClass: db.WorkspaceClass_Default, Type: db.WorkspaceType_Regular, UsageAttributionID: db.NewTeamAttributionID(uuid.New().String()), StartedTime: db.NewVarCharTime(now.Add(1 * time.Minute)), } inserts, updates, err := reconcileUsage([]db.WorkspaceInstanceForUsage{instance, instance}, nil, pricer, now) require.NoError(t, err) require.Len(t, inserts, 1) require.Len(t, updates, 0) expectedUsage := db.Usage{ ID: inserts[0].ID, AttributionID: instance.UsageAttributionID, Description: usageDescriptionFromController, CreditCents: db.NewCreditCents(pricer.CreditsUsedByInstance(&instance, now)), EffectiveTime: db.NewVarCharTime(now), Kind: db.WorkspaceInstanceUsageKind, WorkspaceInstanceID: &instance.ID, Draft: true, Metadata: nil, } require.NoError(t, expectedUsage.SetMetadataWithWorkspaceInstance(db.WorkspaceInstanceUsageData{ WorkspaceId: instance.WorkspaceID, WorkspaceType: instance.Type, WorkspaceClass: instance.WorkspaceClass, ContextURL: instance.ContextURL, StartTime: db.TimeToISO8601(instance.StartedTime.Time()), EndTime: "", UserName: instance.UserName, UserAvatarURL: instance.UserAvatarURL, })) require.EqualValues(t, expectedUsage, inserts[0]) }) t.Run("updates a usage record when a draft exists", func(t *testing.T) { instance := db.WorkspaceInstanceForUsage{ ID: uuid.New(), WorkspaceID: dbtest.GenerateWorkspaceID(), OwnerID: uuid.New(), ProjectID: sql.NullString{ String: "my-project", Valid: true, }, WorkspaceClass: db.WorkspaceClass_Default, Type: db.WorkspaceType_Regular, UsageAttributionID: db.NewTeamAttributionID(uuid.New().String()), StartedTime: db.NewVarCharTime(now.Add(1 * time.Minute)), } // the fields in the usage record deliberately do not match the instance, except for the Instance ID. // we do this to test that the fields in the usage records get updated to reflect the true values from the source of truth - instances. draft := dbtest.NewUsage(t, db.Usage{ ID: uuid.New(), AttributionID: db.NewUserAttributionID(uuid.New().String()), Description: "Some description", CreditCents: 1, EffectiveTime: db.VarcharTime{}, Kind: db.WorkspaceInstanceUsageKind, WorkspaceInstanceID: &instance.ID, Draft: true, Metadata: nil, }) inserts, updates, err := reconcileUsage([]db.WorkspaceInstanceForUsage{instance}, []db.Usage{draft}, pricer, now) require.NoError(t, err) require.Len(t, inserts, 0) require.Len(t, updates, 1) expectedUsage := db.Usage{ ID: draft.ID, AttributionID: instance.UsageAttributionID, Description: usageDescriptionFromController, CreditCents: db.NewCreditCents(pricer.CreditsUsedByInstance(&instance, now)), EffectiveTime: db.NewVarCharTime(now), Kind: db.WorkspaceInstanceUsageKind, WorkspaceInstanceID: &instance.ID, Draft: true, Metadata: nil, } require.NoError(t, expectedUsage.SetMetadataWithWorkspaceInstance(db.WorkspaceInstanceUsageData{ WorkspaceId: instance.WorkspaceID, WorkspaceType: instance.Type, WorkspaceClass: instance.WorkspaceClass, ContextURL: instance.ContextURL, StartTime: db.TimeToISO8601(instance.StartedTime.Time()), EndTime: "", UserName: instance.UserName, UserAvatarURL: instance.UserAvatarURL, })) require.EqualValues(t, expectedUsage, updates[0]) }) t.Run("handles instances without stopping but stopped time", func(t *testing.T) { instance := db.WorkspaceInstanceForUsage{ ID: uuid.New(), WorkspaceID: dbtest.GenerateWorkspaceID(), OwnerID: uuid.New(), ProjectID: sql.NullString{ String: "my-project", Valid: true, }, WorkspaceClass: db.WorkspaceClass_Default, Type: db.WorkspaceType_Regular, UsageAttributionID: db.NewTeamAttributionID(uuid.New().String()), StartedTime: db.NewVarCharTime(now.Add(1 * time.Minute)), StoppedTime: db.NewVarCharTime(now.Add(2 * time.Minute)), } inserts, updates, err := reconcileUsage([]db.WorkspaceInstanceForUsage{instance}, []db.Usage{}, pricer, now) require.NoError(t, err) require.Len(t, inserts, 1) require.Len(t, updates, 0) require.EqualValues(t, db.NewCreditCents(0.17), inserts[0].CreditCents) require.EqualValues(t, instance.StoppedTime, inserts[0].EffectiveTime) }) } func TestGetAndSetCostCenter(t *testing.T) { conn := dbtest.ConnectForTests(t) costCenterUpdates := []*v1.CostCenter{ { AttributionId: string(db.NewUserAttributionID(uuid.New().String())), SpendingLimit: 8000, BillingStrategy: v1.CostCenter_BILLING_STRATEGY_STRIPE, }, { AttributionId: string(db.NewUserAttributionID(uuid.New().String())), SpendingLimit: 500, BillingStrategy: v1.CostCenter_BILLING_STRATEGY_OTHER, }, { AttributionId: string(db.NewTeamAttributionID(uuid.New().String())), SpendingLimit: 8000, BillingStrategy: v1.CostCenter_BILLING_STRATEGY_STRIPE, }, { AttributionId: string(db.NewTeamAttributionID(uuid.New().String())), SpendingLimit: 0, BillingStrategy: v1.CostCenter_BILLING_STRATEGY_OTHER, }, } service := newUsageService(t, conn) for _, costCenter := range costCenterUpdates { retrieved, err := service.SetCostCenter(context.Background(), &v1.SetCostCenterRequest{ CostCenter: costCenter, }) require.NoError(t, err) require.Equal(t, costCenter.SpendingLimit, retrieved.CostCenter.SpendingLimit) require.Equal(t, costCenter.BillingStrategy, retrieved.CostCenter.BillingStrategy) } } func TestListUsage(t *testing.T) { conn := dbtest.ConnectForTests(t) start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC) end := time.Date(2022, 8, 1, 0, 0, 0, 0, time.UTC) attributionID := db.NewTeamAttributionID(uuid.New().String()) draftBefore := dbtest.NewUsage(t, db.Usage{ AttributionID: attributionID, EffectiveTime: db.NewVarCharTime(start.Add(-1 * 23 * time.Hour)), CreditCents: 100, Draft: true, }) nondraftBefore := dbtest.NewUsage(t, db.Usage{ AttributionID: attributionID, EffectiveTime: db.NewVarCharTime(start.Add(-1 * 23 * time.Hour)), CreditCents: 200, Draft: false, }) draftInside := dbtest.NewUsage(t, db.Usage{ AttributionID: attributionID, EffectiveTime: db.NewVarCharTime(start.Add(2 * time.Hour)), CreditCents: 300, Draft: true, }) nonDraftInside := dbtest.NewUsage(t, db.Usage{ AttributionID: attributionID, EffectiveTime: db.NewVarCharTime(start.Add(2 * time.Hour)), CreditCents: 400, Draft: false, }) nonDraftAfter := dbtest.NewUsage(t, db.Usage{ AttributionID: attributionID, EffectiveTime: db.NewVarCharTime(end.Add(2 * time.Hour)), CreditCents: 1000, }) dbtest.CreateUsageRecords(t, conn, draftBefore, nondraftBefore, draftInside, nonDraftInside, nonDraftAfter) usageService := newUsageService(t, conn) tests := []struct { start, end time.Time // expectations creditsUsed float64 recordsInRange int64 }{ {start, end, 7, 2}, {end, end, 0, 0}, {start, start, 0, 0}, {start.Add(-200 * 24 * time.Hour), end, 10, 4}, {start.Add(-200 * 24 * time.Hour), end.Add(10 * 24 * time.Hour), 20, 5}, } for i, test := range tests { t.Run(fmt.Sprintf("Running test no %d", i+1), func(t *testing.T) { metaData, err := usageService.ListUsage(context.Background(), &v1.ListUsageRequest{ AttributionId: string(attributionID), From: timestamppb.New(test.start), To: timestamppb.New(test.end), Order: v1.ListUsageRequest_ORDERING_DESCENDING, Pagination: &v1.PaginatedRequest{ PerPage: 1, Page: 1, }, }) require.NoError(t, err) require.Equal(t, test.creditsUsed, metaData.CreditsUsed) require.Equal(t, test.recordsInRange, metaData.Pagination.Total) }) } } func TestAddUSageCreditNote(t *testing.T) { conn := dbtest.ConnectForTests(t) attributionID := db.NewTeamAttributionID(uuid.New().String()) usageService := newUsageService(t, conn) tests := []struct { credits int32 userId string description string // expectations expectedError bool }{ {300, uuid.New().String(), "Something", false}, {300, "bad-userid", "Something", true}, {300, uuid.New().String(), " " /* no note */, true}, {-300, uuid.New().String(), "Negative Balance", false}, } for i, test := range tests { t.Run(fmt.Sprintf("Running test no %d", i+1), func(t *testing.T) { _, err := usageService.AddUsageCreditNote(context.Background(), &v1.AddUsageCreditNoteRequest{ AttributionId: string(attributionID), Credits: test.credits, Description: test.description, UserId: test.userId, }) if test.expectedError { require.Error(t, err) } else { require.NoError(t, err) balance, err := db.GetBalance(context.Background(), conn, attributionID) require.NoError(t, err) require.Equal(t, int32(balance.ToCredits()), test.credits*-1) } require.NoError(t, conn.Where("attributionId = ?", attributionID).Delete(&db.Usage{}).Error) }) } }