feat: fix and add more tests

Author: karlem

core/vesting/engine.go

@@ -68,8 +68,8 @@ type PartyRewards struct {
 }
 
 type multiplierAndQuantBalance struct {
-	multiplier     num.Decimal
-	quantumBalance num.Decimal
+	Multiplier     num.Decimal
+	QuantumBalance num.Decimal
 }
 
 type Engine struct {
@@ -202,32 +202,32 @@ func (e *Engine) getSummedBonusMultiplier(keys []string) multiplierAndQuantBalan
 			continue
 		}
 
-		summed.quantumBalance = summed.quantumBalance.Add(quantumBalanceForKey)
+		summed.QuantumBalance = summed.QuantumBalance.Add(quantumBalanceForKey)
 	}
 
-	summed.multiplier = e.rewardBonusMultiplier(summed.quantumBalance)
+	summed.Multiplier = e.rewardBonusMultiplier(summed.QuantumBalance)
 
 	return summed
 }
 
 // GetRewardBonusMultiplier returns the reward bonus multiplier and quantum balance sum for a party and it's derived keys.
 func (e *Engine) GetRewardBonusMultiplier(party string) (num.Decimal, num.Decimal) {
-	partyID, derivedKeys := e.parties.RelatedKeys(party)
-
 	owner := party
+
+	partyID, derivedKeys := e.parties.RelatedKeys(party)
 	if partyID != nil {
 		owner = partyID.String()
 	}
 
 	if cached, ok := e.summedRewardBonusMultiplierCache[owner]; ok {
-		return cached.quantumBalance, cached.multiplier
+		return cached.QuantumBalance, cached.Multiplier
 	}
 
 	summed := e.getSummedBonusMultiplier(append(derivedKeys, owner))
 
 	e.summedRewardBonusMultiplierCache[owner] = summed
 
-	return summed.quantumBalance, summed.multiplier
+	return summed.QuantumBalance, summed.Multiplier
 }
 
 // GetSingleAndSummedRewardBonusMultipliers returns a single and summed reward bonus multipliers and quantum balances for a party.
@@ -238,13 +238,15 @@ func (e *Engine) GetSingleAndSummedRewardBonusMultipliers(party string) (
 	single multiplierAndQuantBalance,
 	summed multiplierAndQuantBalance,
 ) {
+	owner := party
+
 	partyID, derivedKeys := e.parties.RelatedKeys(party)
+	if partyID != nil {
+		owner = partyID.String()
+	}
 
 	single = e.getSummedBonusMultiplier([]string{party})
 
-	// calculate with derived keys. First add the owner's quantum balance.
-	owner := partyID.String()
-
 	if cached, ok := e.summedRewardBonusMultiplierCache[owner]; ok {
 		return single, cached
 	}
@@ -495,14 +497,14 @@ func (e *Engine) broadcastVestingStatsUpdate(ctx context.Context, seq uint64) {
 	for _, party := range parties {
 		single, summed := e.GetSingleAndSummedRewardBonusMultipliers(party)
 		// To avoid excessively large decimals.
-		single.quantumBalance.Round(2)
-		summed.quantumBalance.Round(2)
+		single.QuantumBalance.Round(2)
+		summed.QuantumBalance.Round(2)
 		evt.Stats = append(evt.Stats, &eventspb.PartyVestingStats{
 			PartyId:                     party,
-			RewardBonusMultiplier:       single.multiplier.String(),
-			QuantumBalance:              single.quantumBalance.String(),
-			SummedRewardBonusMultiplier: summed.multiplier.String(),
-			SummedQuantumBalance:        summed.quantumBalance.String(),
+			RewardBonusMultiplier:       single.Multiplier.String(),
+			QuantumBalance:              single.QuantumBalance.String(),
+			SummedRewardBonusMultiplier: summed.Multiplier.String(),
+			SummedQuantumBalance:        summed.QuantumBalance.String(),
 		})
 	}
 

core/vesting/engine_test.go

@@ -1069,21 +1069,21 @@ func TestDistributeWithRelatedKeys(t *testing.T) {
 				Stats: []*eventspb.PartyVestingStats{
 					{
 						PartyId:                     derivedKeys[0],
-						RewardBonusMultiplier:       "2",
+						RewardBonusMultiplier:       "1",
 						QuantumBalance:              "145",
 						SummedRewardBonusMultiplier: "3",
 						SummedQuantumBalance:        "825",
 					},
 					{
 						PartyId:                     derivedKeys[1],
-						RewardBonusMultiplier:       "2",
+						RewardBonusMultiplier:       "1",
 						QuantumBalance:              "145",
 						SummedRewardBonusMultiplier: "3",
 						SummedQuantumBalance:        "825",
 					},
 					{
 						PartyId:                     derivedKeys[2],
-						RewardBonusMultiplier:       "2",
+						RewardBonusMultiplier:       "1",
 						QuantumBalance:              "145",
 						SummedRewardBonusMultiplier: "3",
 						SummedQuantumBalance:        "825",
@@ -1166,12 +1166,33 @@ func TestDistributeWithRelatedKeys(t *testing.T) {
 			assert.Equal(t, eventspb.VestingStatsUpdated{
 				AtEpoch: epochSeq,
 				Stats: []*eventspb.PartyVestingStats{
+					{
+						PartyId:                     derivedKeys[0],
+						RewardBonusMultiplier:       "1",
+						QuantumBalance:              "150",
+						SummedRewardBonusMultiplier: "3",
+						SummedQuantumBalance:        "850",
+					},
+					{
+						PartyId:                     derivedKeys[1],
+						RewardBonusMultiplier:       "1",
+						QuantumBalance:              "150",
+						SummedRewardBonusMultiplier: "3",
+						SummedQuantumBalance:        "850",
+					},
+					{
+						PartyId:                     derivedKeys[2],
+						RewardBonusMultiplier:       "1",
+						QuantumBalance:              "150",
+						SummedRewardBonusMultiplier: "3",
+						SummedQuantumBalance:        "850",
+					},
 					{
 						PartyId:                     party,
 						RewardBonusMultiplier:       "2",
 						QuantumBalance:              "400",
-						SummedRewardBonusMultiplier: "2",
-						SummedQuantumBalance:        "400",
+						SummedRewardBonusMultiplier: "3",
+						SummedQuantumBalance:        "850",
 					},
 				},
 			}, e.Proto())
@@ -1219,6 +1240,116 @@ func TestDistributeWithRelatedKeys(t *testing.T) {
 	})
 }
 
+func TestGetRewardBonusMultiplier(t *testing.T) {
+	v := getTestEngine(t)
+
+	ctx := context.Background()
+
+	// distribute 90% as the base rate,
+	// so first we distribute some, then we get under the minimum value, and all the rest
+	// is distributed
+	require.NoError(t, v.OnRewardVestingBaseRateUpdate(ctx, num.MustDecimalFromString("0.9")))
+	// this is multiplied by the quantum, so it will make it 100% of the quantum
+	require.NoError(t, v.OnRewardVestingMinimumTransferUpdate(ctx, num.MustDecimalFromString("1")))
+
+	require.NoError(t, v.OnBenefitTiersUpdate(ctx, &vegapb.VestingBenefitTiers{
+		Tiers: []*vegapb.VestingBenefitTier{
+			{
+				MinimumQuantumBalance: "200",
+				RewardMultiplier:      "1",
+			},
+			{
+				MinimumQuantumBalance: "500",
+				RewardMultiplier:      "2",
+			},
+			{
+				MinimumQuantumBalance: "1200",
+				RewardMultiplier:      "3",
+			},
+		},
+	}))
+
+	party := "party1"
+	partyID := types.PartyID(party)
+	vegaAsset := "VEGA"
+	derivedKeys := []string{"derived1", "derived2", "derived3", "derived4"}
+
+	v.parties.EXPECT().RelatedKeys(party).Return(&partyID, derivedKeys).AnyTimes()
+
+	v.col.InitVestedBalance(party, vegaAsset, num.NewUint(500))
+
+	for _, key := range derivedKeys {
+		v.col.InitVestedBalance(key, vegaAsset, num.NewUint(250))
+		v.parties.EXPECT().RelatedKeys(key).Return(&partyID, derivedKeys).AnyTimes()
+	}
+
+	for _, key := range append(derivedKeys, party) {
+		balance, multiplier := v.GetRewardBonusMultiplier(key)
+		require.Equal(t, num.DecimalFromInt64(1500), balance)
+		require.Equal(t, num.DecimalFromInt64(3), multiplier)
+	}
+
+	// check that we only called the GetVestingQuantumBalance once for each key
+	// later calls should be cached
+	require.Equal(t, 5, v.col.GetVestingQuantumBalanceCallCount())
+
+	v.col.ResetVestingQuantumBalanceCallCount()
+
+	single, summed := v.GetSingleAndSummedRewardBonusMultipliers(party)
+	require.Equal(t, num.DecimalFromInt64(500), single.QuantumBalance)
+	require.Equal(t, num.DecimalFromInt64(2), single.Multiplier)
+	require.Equal(t, num.DecimalFromInt64(1500), summed.QuantumBalance)
+	require.Equal(t, num.DecimalFromInt64(3), summed.Multiplier)
+
+	for _, key := range derivedKeys {
+		single, summed := v.GetSingleAndSummedRewardBonusMultipliers(key)
+		require.Equal(t, num.DecimalFromInt64(250), single.QuantumBalance)
+		require.Equal(t, num.DecimalFromInt64(1), single.Multiplier)
+		require.Equal(t, num.DecimalFromInt64(1500), summed.QuantumBalance)
+		require.Equal(t, num.DecimalFromInt64(3), summed.Multiplier)
+	}
+
+	// here we check that we just called 5 more times to re-calculate single values
+	// summed should be cached
+	require.Equal(t, 5, v.col.GetVestingQuantumBalanceCallCount())
+
+	v.broker.EXPECT().Send(gomock.Any()).AnyTimes()
+
+	// now we simulate the end of the epoch
+	// it will reset cache for reward bonus multipliers
+	v.OnEpochEvent(ctx, types.Epoch{
+		Action: vegapb.EpochAction_EPOCH_ACTION_END,
+		Seq:    1,
+	})
+
+	v.col.ResetVestingQuantumBalanceCallCount()
+
+	for _, key := range append(derivedKeys, party) {
+		balance, multiplier := v.GetRewardBonusMultiplier(key)
+		require.Equal(t, num.DecimalFromInt64(1500), balance)
+		require.Equal(t, num.DecimalFromInt64(3), multiplier)
+	}
+
+	// now it's called 5 times again because the cache gets reset at the end of the epoch
+	require.Equal(t, 5, v.col.GetVestingQuantumBalanceCallCount())
+
+	v.OnEpochEvent(ctx, types.Epoch{
+		Action: vegapb.EpochAction_EPOCH_ACTION_END,
+		Seq:    1,
+	})
+
+	v.col.ResetVestingQuantumBalanceCallCount()
+
+	for _, key := range append(derivedKeys, party) {
+		balance, multiplier := v.GetRewardBonusMultiplier(key)
+		require.Equal(t, num.DecimalFromInt64(1500), balance)
+		require.Equal(t, num.DecimalFromInt64(3), multiplier)
+	}
+
+	// now it's called 5 times again because the cache gets reset at the end of the epoch
+	require.Equal(t, 5, v.col.GetVestingQuantumBalanceCallCount())
+}
+
 // LedgerMovements is the result of a mock, so it doesn't really make sense to
 // verify data consistency.
 func expectLedgerMovements(t *testing.T, v *testEngine) {

core/vesting/helper_for_test.go

@@ -108,7 +108,8 @@ func newEngine(t *testing.T) *testSnapshotEngine {
 }
 
 type collateralMock struct {
-	vestedAccountAmount map[string]map[string]*num.Uint
+	vestedAccountAmount            map[string]map[string]*num.Uint
+	vestingQuantumBalanceCallCount int
 }
 
 func (c *collateralMock) InitVestedBalance(party, asset string, balance *num.Uint) {
@@ -156,9 +157,19 @@ func (c *collateralMock) GetAllVestingQuantumBalance(party string) num.Decimal {
 		balance = balance.Add(num.DecimalFromUint(n))
 	}
 
+	c.vestingQuantumBalanceCallCount += 1
+
 	return balance
 }
 
+func (c *collateralMock) ResetVestingQuantumBalanceCallCount() {
+	c.vestingQuantumBalanceCallCount = 0
+}
+
+func (c *collateralMock) GetVestingQuantumBalanceCallCount() int {
+	return c.vestingQuantumBalanceCallCount
+}
+
 func newCollateralMock(t *testing.T) *collateralMock {
 	t.Helper()