refactor(balances): return features that will have balance in next wi…

…ndow

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

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