Implement multi-block PSE community distribution

integration-tests/modules/pse_test.go

@@ -251,9 +251,7 @@ func TestPSEDistribution(t *testing.T) {
 	requireT.NoError(err)
 	t.Logf("Distribution 1 at height: %d", height)
 
-	scheduledDistributions, err := getScheduledDistribution(ctx, chain)
-	requireT.NoError(err)
-	requireT.Len(scheduledDistributions, 2)
+	awaitScheduleCount(ctx, t, chain, 2)
 
 	balancesBefore, scoresBefore, totalScore := getAllDelegatorInfo(ctx, t, chain, height-1)
 	balancesAfter, _, _ := getAllDelegatorInfo(ctx, t, chain, height)
@@ -298,9 +296,7 @@ func TestPSEDistribution(t *testing.T) {
 	requireT.NoError(err)
 	t.Logf("Distribution 2 at height: %d", height)
 
-	scheduledDistributions, err = getScheduledDistribution(ctx, chain)
-	requireT.NoError(err)
-	requireT.Len(scheduledDistributions, 1)
+	awaitScheduleCount(ctx, t, chain, 1)
 
 	balancesBefore, scoresBefore, totalScore = getAllDelegatorInfo(ctx, t, chain, height-1)
 	balancesAfter, _, _ = getAllDelegatorInfo(ctx, t, chain, height)
@@ -336,9 +332,7 @@ func TestPSEDistribution(t *testing.T) {
 	requireT.NoError(err)
 	t.Logf("Distribution 3 at height: %d", height)
 
-	scheduledDistributions, err = getScheduledDistribution(ctx, chain)
-	requireT.NoError(err)
-	requireT.Empty(scheduledDistributions)
+	awaitScheduleCount(ctx, t, chain, 0)
 
 	balancesBefore, scoresBefore, totalScore = getAllDelegatorInfo(ctx, t, chain, height-1)
 	balancesAfter, _, _ = getAllDelegatorInfo(ctx, t, chain, height)
@@ -836,7 +830,7 @@ func awaitScheduledDistributionEvent(
 ) (int64, communityDistributedEvent, error) {
 	var observedHeight int64
 	err := chain.AwaitState(ctx, func(ctx context.Context) error {
-		query := fmt.Sprintf("tx.pse.v1.EventAllocationDistributed.mode='EndBlock' AND block.height>%d", startHeight)
+		query := fmt.Sprintf("tx.pse.v1.EventCommunityDistributed.mode='EndBlock' AND block.height>%d", startHeight)
 		blocks, err := chain.ClientContext.RPCClient().BlockSearch(ctx, query, nil, nil, "")
 		if err != nil {
 			return err
@@ -880,6 +874,22 @@ func getScheduledDistribution(
 	return pseResponse.ScheduledDistributions, nil
 }
 
+func awaitScheduleCount(ctx context.Context, t *testing.T, chain integration.TXChain, expectedCount int) {
+	t.Helper()
+	requireT := require.New(t)
+	err := chain.AwaitState(ctx, func(ctx context.Context) error {
+		dist, err := getScheduledDistribution(ctx, chain)
+		if err != nil {
+			return err
+		}
+		if len(dist) != expectedCount {
+			return fmt.Errorf("expected %d scheduled distributions, got %d", expectedCount, len(dist))
+		}
+		return nil
+	}, integration.WithAwaitStateTimeout(10*time.Second))
+	requireT.NoError(err)
+}
+
 func removeAttributeFromEvent(events []tmtypes.Event, key string) []tmtypes.Event {
 	newEvents := make([]tmtypes.Event, 0, len(events))
 	for _, event := range events {

x/pse/keeper/delegation.go

@@ -57,10 +57,7 @@ func (k Keeper) GetDelegatorScore(
 
 // SetDelegatorScore sets the score for a delegator.
 func (k Keeper) SetDelegatorScore(
-	ctx context.Context,
-	distributionID uint64,
-	delAddr sdk.AccAddress,
-	score sdkmath.Int,
+	ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int,
 ) error {
 	key := collections.Join(distributionID, delAddr)
 	return k.AccountScoreSnapshot.Set(ctx, key, score)
@@ -72,6 +69,22 @@ func (k Keeper) RemoveDelegatorScore(ctx context.Context, distributionID uint64,
 	return k.AccountScoreSnapshot.Remove(ctx, key)
 }
 
+// addToScore atomically adds a score value to a delegator's score snapshot.
+func (k Keeper) addToScore(
+	ctx context.Context, distributionID uint64, delAddr sdk.AccAddress, score sdkmath.Int,
+) error {
+	if score.IsZero() {
+		return nil
+	}
+	lastScore, err := k.GetDelegatorScore(ctx, distributionID, delAddr)
+	if errors.Is(err, collections.ErrNotFound) {
+		lastScore = sdkmath.NewInt(0)
+	} else if err != nil {
+		return err
+	}
+	return k.SetDelegatorScore(ctx, distributionID, delAddr, lastScore.Add(score))
+}
+
 // CalculateDelegatorScore calculates the current total score for a delegator.
 // This includes both the accumulated score snapshot (from previous periods)
 // and the current period score calculated on-demand from active delegations.
@@ -82,7 +95,7 @@ func (k Keeper) CalculateDelegatorScore(ctx context.Context, delAddr sdk.AccAddr
 	if err != nil {
 		return sdkmath.Int{}, err
 	}
-	distributionID := distribution.ID // TODO update to handle distribution ID properly.
+	distributionID := distribution.ID
 
 	// Start with the accumulated score from the snapshot (previous periods)
 	accumulatedScore, err := k.GetDelegatorScore(ctx, distributionID, delAddr)