diff --git a/internal/credit/postgres_connector/balance.go b/internal/credit/postgres_connector/balance.go index 0c7fa3c21..f324c8bf9 100644 --- a/internal/credit/postgres_connector/balance.go +++ b/internal/credit/postgres_connector/balance.go @@ -74,22 +74,37 @@ func (a *PostgresConnector) getBalance( // Get features in grants features := map[credit.FeatureID]credit.Feature{} - for _, grant := range grants { - if grant.Void { - ledgerEntries.AddVoidGrant(grant) - continue + { + // We should displays features that'll have active balances in the next window + // to avoid confusion + toPlusWindow := to.Add(a.config.WindowSize) + grants, err := a.ListGrants(ctx, credit.ListGrantsParams{ + Namespace: ledger.Namespace, + LedgerIDs: []credit.LedgerID{ledgerID.ID}, + From: &from, + To: &toPlusWindow, + IncludeVoid: true, + }) + if err != nil { + return credit.Balance{}, ledgerEntries, fmt.Errorf("list grants: %w", err) } + for _, grant := range grants { + if grant.Void { + ledgerEntries.AddVoidGrant(grant) + continue + } - ledgerEntries.AddGrant(grant) + ledgerEntries.AddGrant(grant) - if grant.FeatureID != nil { - featureID := *grant.FeatureID - if _, ok := features[featureID]; !ok { - feature, err := a.GetFeature(ctx, credit.NewNamespacedFeatureID(ledgerID.Namespace, featureID)) - if err != nil { - return credit.Balance{}, ledgerEntries, fmt.Errorf("get feature: %w", err) + if grant.FeatureID != nil { + featureID := *grant.FeatureID + if _, ok := features[featureID]; !ok { + feature, err := a.GetFeature(ctx, credit.NewNamespacedFeatureID(ledgerID.Namespace, featureID)) + if err != nil { + return credit.Balance{}, ledgerEntries, fmt.Errorf("get feature: %w", err) + } + features[featureID] = feature } - features[featureID] = feature } } } @@ -298,6 +313,14 @@ func (a *PostgresConnector) getBalance( // Aggregate grant balances by feature featureBalancesMap := map[credit.FeatureID]credit.FeatureBalance{} + for _, feature := range features { + featureBalancesMap[*feature.ID] = credit.FeatureBalance{ + Feature: feature, + Balance: 0, + Usage: 0, + } + + } for _, grantBalance := range grantBalances { if grantBalance.FeatureID == nil { continue diff --git a/internal/credit/postgres_connector/balance_test.go b/internal/credit/postgres_connector/balance_test.go index d28142715..e01c5ebd7 100644 --- a/internal/credit/postgres_connector/balance_test.go +++ b/internal/credit/postgres_connector/balance_test.go @@ -1060,6 +1060,88 @@ func TestPostgresConnectorBalances(t *testing.T) { assert.Equal(t, grant2.Amount, grant2FromBalance.Balance) }, }, + { + name: "Should return features the subjects granted to in the near future up to windowsize, even if current balance & total are 0", + description: `We should return all features they have grants to, even if the balance is 0 due to no active grants`, + test: func(t *testing.T, connector credit.Connector, streamingConnector *testutils.MockStreamingConnector, db_client *db.Client) { + start, err := time.Parse(time.RFC3339, "2024-01-01T00:00:00Z") + assert.NoError(t, err) + subject := ulid.Make().String() + ledgerID := credit.LedgerID(ulid.Make().String()) + + // Create Ledger + + ledger, err := connector.CreateLedger(context.Background(), credit.Ledger{ + Namespace: namespace, + ID: ledgerID, + Subject: subject, + CreatedAt: start, + }) + assert.NoError(t, err) + + now := start.Add(windowSize) + future := start.Add(windowSize * 2) + + // Create Feature + feature1 := testutils.CreateFeature(t, connector, featureIn1) + feature2 := testutils.CreateFeature(t, connector, featureIn2) + // Meters need events so they're found + streamingConnector.AddSimpleEvent(meter1.Slug, 1, now.Add(time.Hour*1000)) + streamingConnector.AddSimpleEvent(meter2.Slug, 1, now.Add(time.Hour*1000)) + // Create grants for both + _, err = connector.CreateGrant(context.Background(), credit.Grant{ + Namespace: namespace, + LedgerID: ledger.ID, + FeatureID: feature1.ID, + Type: credit.GrantTypeUsage, + Amount: 100, + Priority: 1, + EffectiveAt: future, + Expiration: credit.ExpirationPeriod{ + Duration: credit.ExpirationPeriodDurationYear, + Count: 1, + }, + CreatedAt: &now, + UpdatedAt: &now, + }) + assert.NoError(t, err) + + _, err = connector.CreateGrant(context.Background(), credit.Grant{ + Namespace: namespace, + LedgerID: ledger.ID, + FeatureID: feature2.ID, + Type: credit.GrantTypeUsage, + Amount: 100, + Priority: 0, + EffectiveAt: future.Add(windowSize), + Expiration: credit.ExpirationPeriod{ + Duration: credit.ExpirationPeriodDurationYear, + Count: 1, + }, + CreatedAt: &now, + UpdatedAt: &now, + }) + assert.NoError(t, err) + + // Get Balance + balance, err := connector.GetBalance(context.Background(), credit.NewNamespacedLedgerID(namespace, ledger.ID), now.Add(time.Minute)) + assert.NoError(t, err) + + // Assert balance shows first feature but not second + assert.Equal(t, 1, len(balance.FeatureBalances)) + assert.ElementsMatch(t, + testutils.RemoveTimestampsFromFeatureBalances([]credit.FeatureBalance{ + { + Feature: feature1, + Balance: 0, + Usage: 0, + }, + }), + testutils.RemoveTimestampsFromFeatureBalances(balance.FeatureBalances), + ) + + }, + }, } for _, tc := range tt { t.Run(tc.name, func(t *testing.T) {