From f5b11dc29b51440095ca8558699c95503b42e7e3 Mon Sep 17 00:00:00 2001
From: beer-1 <147697694+beer-1@users.noreply.github.com>
Date: Fri, 21 Jun 2024 14:40:22 +0900
Subject: [PATCH 1/5] minor correction to aviod incorrect calculation

---
 x/distribution/keeper/keeper.go    |  1 +
 x/distribution/keeper/validator.go |  3 +-
 x/distribution/types/dec_pool.go   | 91 +++++++++++++++---------------
 x/distribution/types/pool.go       | 85 ++++++++++++++--------------
 4 files changed, 90 insertions(+), 90 deletions(-)

diff --git a/x/distribution/keeper/keeper.go b/x/distribution/keeper/keeper.go
index d8a27f3b..64912e9e 100644
--- a/x/distribution/keeper/keeper.go
+++ b/x/distribution/keeper/keeper.go
@@ -178,6 +178,7 @@ func (k Keeper) WithdrawValidatorCommission(ctx context.Context, valAddr sdk.Val
 	}
 
 	commissions, remainder := accumCommission.Commissions.TruncateDecimal()
+
 	// leave remainder to withdraw later
 	if err = k.ValidatorAccumulatedCommissions.Set(ctx, valAddr, customtypes.ValidatorAccumulatedCommission{Commissions: remainder}); err != nil {
 		return nil, err
diff --git a/x/distribution/keeper/validator.go b/x/distribution/keeper/validator.go
index 72295ea2..06f4aecb 100644
--- a/x/distribution/keeper/validator.go
+++ b/x/distribution/keeper/validator.go
@@ -72,9 +72,8 @@ func (k Keeper) IncrementValidatorPeriod(ctx context.Context, val stakingtypes.V
 			// can't calculate ratio for zero-token validators
 			// ergo we instead add to the community pool
 			communityFunding = communityFunding.Add(customtypes.NewDecPool(token.Denom, rewardCoins))
-			current = append(current, customtypes.NewDecPool(token.Denom, sdk.DecCoins{}))
 		} else {
-			current = append(current, customtypes.NewDecPool(token.Denom, rewardCoins.QuoDecTruncate(math.LegacyNewDecFromInt(token.Amount))))
+			current = current.Add(customtypes.NewDecPool(token.Denom, rewardCoins.QuoDecTruncate(math.LegacyNewDecFromInt(token.Amount))))
 		}
 	}
 
diff --git a/x/distribution/types/dec_pool.go b/x/distribution/types/dec_pool.go
index de451d6d..dc54c98d 100644
--- a/x/distribution/types/dec_pool.go
+++ b/x/distribution/types/dec_pool.go
@@ -3,7 +3,6 @@ package types
 import (
 	"fmt"
 	"sort"
-	"strings"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -12,11 +11,18 @@ import (
 // rewards pools for multi-token staking
 type DecPools []DecPool
 
+// NewDecPools creates a new DecPools instance
+func NewDecPools(pools ...DecPool) DecPools {
+	return removeZeroDecPools(pools)
+}
+
 // NewDecPoolsFromPools create DecPools from Pools
 func NewDecPoolsFromPools(pools Pools) DecPools {
 	decPools := DecPools{}
 	for _, p := range pools {
-		decPools = append(decPools, NewDecPool(p.Denom, sdk.NewDecCoinsFromCoins(p.Coins...)))
+		if !p.IsEmpty() {
+			decPools = append(decPools, NewDecPool(p.Denom, sdk.NewDecCoinsFromCoins(p.Coins...)))
+		}
 	}
 
 	return decPools
@@ -37,52 +43,38 @@ func (pools DecPools) Add(poolsB ...DecPool) DecPools {
 }
 
 // Add will perform addition of two DecPools sets.
-func (pools DecPools) safeAdd(poolsB DecPools) DecPools {
-	sum := ([]DecPool)(nil)
-	indexA, indexB := 0, 0
-	lenA, lenB := len(pools), len(poolsB)
-
-	for {
-		if indexA == lenA {
-			if indexB == lenB {
-				// return nil pools if both sets are empty
-				return sum
-			}
-
-			// return set B (excluding zero pools) if set A is empty
-			return append(sum, removeZeroDecPools(poolsB[indexB:])...)
-		} else if indexB == lenB {
-			// return set A (excluding zero pools) if set B is empty
-			return append(sum, removeZeroDecPools(pools[indexA:])...)
-		}
-
-		poolA, poolB := pools[indexA], poolsB[indexB]
-
-		switch strings.Compare(poolA.Denom, poolB.Denom) {
-		case -1: // pool A denom < pool B denom
-			if !poolA.IsEmpty() {
-				sum = append(sum, poolA)
-			}
-
-			indexA++
-
-		case 0: // pool A denom == pool B denom
-			res := poolA.Add(poolB)
-			if !res.IsEmpty() {
-				sum = append(sum, res)
-			}
-
-			indexA++
-			indexB++
+func (pools DecPools) safeAdd(poolsB DecPools) (coalesced DecPools) {
+	// probably the best way will be to make DecPools and interface and hide the structure
+	// definition (type alias)
+	if !pools.isSorted() {
+		panic("DecPools (self) must be sorted")
+	}
+	if !poolsB.isSorted() {
+		panic("Wrong argument: DecPools must be sorted")
+	}
 
-		case 1: // pool A denom > pool B denom
-			if !poolB.IsEmpty() {
-				sum = append(sum, poolB)
-			}
+	uniqPools := make(map[string]DecPools, len(pools)+len(poolsB))
+	// Traverse all the pools for each of the pools and poolsB.
+	for _, pL := range []DecPools{pools, poolsB} {
+		for _, p := range pL {
+			uniqPools[p.Denom] = append(uniqPools[p.Denom], p)
+		}
+	}
 
-			indexB++
+	for denom, pL := range uniqPools { //#nosec
+		comboPool := DecPool{Denom: denom, DecCoins: sdk.DecCoins{}}
+		for _, p := range pL {
+			comboPool = comboPool.Add(p)
+		}
+		if !comboPool.IsEmpty() {
+			coalesced = append(coalesced, comboPool)
 		}
 	}
+	if coalesced == nil {
+		return DecPools{}
+	}
+
+	return coalesced.Sort()
 }
 
 // Sub subtracts a set of DecPools from another (adds the inverse).
@@ -260,6 +252,15 @@ func (p DecPools) Sort() DecPools {
 	return p
 }
 
+func (p DecPools) isSorted() bool {
+	for i := 1; i < len(p); i++ {
+		if p[i-1].Denom > p[i].Denom {
+			return false
+		}
+	}
+	return true
+}
+
 //-----------------------------------------------------------------------------
 // DecPool functions
 
@@ -297,7 +298,7 @@ func (pool DecPool) Sub(poolB DecPool) DecPool {
 // change. Note, the change may be zero.
 func (pool DecPool) TruncateDecimal() (Pool, DecPool) {
 	truncated, change := pool.DecCoins.TruncateDecimal()
-	return NewPool(pool.Denom, sdk.NewCoins(truncated...)), NewDecPool(pool.Denom, change)
+	return NewPool(pool.Denom, truncated), NewDecPool(pool.Denom, change)
 }
 
 func removeZeroDecPools(pools DecPools) DecPools {
diff --git a/x/distribution/types/pool.go b/x/distribution/types/pool.go
index 65e2862b..bf38ddd6 100644
--- a/x/distribution/types/pool.go
+++ b/x/distribution/types/pool.go
@@ -3,7 +3,6 @@ package types
 import (
 	"fmt"
 	"sort"
-	"strings"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -12,6 +11,11 @@ import (
 // rewards pools for multi-token staking
 type Pools []Pool
 
+// NewPools creates a new Pools instance
+func NewPools(pools ...Pool) Pools {
+	return removeZeroPools(pools)
+}
+
 // Sum returns sum of pool tokens
 func (p Pools) Sum() (coins sdk.Coins) {
 	for _, p := range p {
@@ -27,52 +31,38 @@ func (pools Pools) Add(poolsB ...Pool) Pools {
 }
 
 // Add will perform addition of two Pools sets.
-func (pools Pools) safeAdd(poolsB Pools) Pools {
-	sum := ([]Pool)(nil)
-	indexA, indexB := 0, 0
-	lenA, lenB := len(pools), len(poolsB)
-
-	for {
-		if indexA == lenA {
-			if indexB == lenB {
-				// return nil pools if both sets are empty
-				return sum
-			}
-
-			// return set B (excluding zero pools) if set A is empty
-			return append(sum, removeZeroPools(poolsB[indexB:])...)
-		} else if indexB == lenB {
-			// return set A (excluding zero pools) if set B is empty
-			return append(sum, removeZeroPools(pools[indexA:])...)
-		}
-
-		poolA, poolB := pools[indexA], poolsB[indexB]
-
-		switch strings.Compare(poolA.Denom, poolB.Denom) {
-		case -1: // pool A denom < pool B denom
-			if !poolA.IsEmpty() {
-				sum = append(sum, poolA)
-			}
-
-			indexA++
-
-		case 0: // pool A denom == pool B denom
-			res := poolA.Add(poolB)
-			if !res.IsEmpty() {
-				sum = append(sum, res)
-			}
-
-			indexA++
-			indexB++
+func (pools Pools) safeAdd(poolsB Pools) (coalesced Pools) {
+	// probably the best way will be to make Pools and interface and hide the structure
+	// definition (type alias)
+	if !pools.isSorted() {
+		panic("Pools (self) must be sorted")
+	}
+	if !poolsB.isSorted() {
+		panic("Wrong argument: Pools must be sorted")
+	}
 
-		case 1: // pool A denom > pool B denom
-			if !poolB.IsEmpty() {
-				sum = append(sum, poolB)
-			}
+	uniqPools := make(map[string]Pools, len(pools)+len(poolsB))
+	// Traverse all the pools for each of the pools and poolsB.
+	for _, pL := range []Pools{pools, poolsB} {
+		for _, c := range pL {
+			uniqPools[c.Denom] = append(uniqPools[c.Denom], c)
+		}
+	}
 
-			indexB++
+	for denom, pL := range uniqPools { //#nosec
+		comboPool := Pool{Denom: denom, Coins: sdk.Coins{}}
+		for _, p := range pL {
+			comboPool = comboPool.Add(p)
+		}
+		if !comboPool.IsEmpty() {
+			coalesced = append(coalesced, comboPool)
 		}
 	}
+	if coalesced == nil {
+		return Pools{}
+	}
+
+	return coalesced.Sort()
 }
 
 // Sub subtracts a set of Pools from another (adds the inverse).
@@ -216,6 +206,15 @@ func (p Pools) Sort() Pools {
 	return p
 }
 
+func (p Pools) isSorted() bool {
+	for i := 1; i < len(p); i++ {
+		if p[i-1].Denom > p[i].Denom {
+			return false
+		}
+	}
+	return true
+}
+
 //-----------------------------------------------------------------------------
 // Pool functions
 

From c51dfe7cbe69735c31e258ec0a4407b087067fff Mon Sep 17 00:00:00 2001
From: sh-cha <snghuncha@gmail.com>
Date: Tue, 25 Jun 2024 16:22:59 +0900
Subject: [PATCH 2/5] add pool test & fix something that could be logically
 wrong

---
 x/distribution/types/pool.go      |  25 ++-
 x/distribution/types/pool_test.go | 268 ++++++++++++++++++++++++++++++
 2 files changed, 284 insertions(+), 9 deletions(-)
 create mode 100644 x/distribution/types/pool_test.go

diff --git a/x/distribution/types/pool.go b/x/distribution/types/pool.go
index bf38ddd6..67f333c4 100644
--- a/x/distribution/types/pool.go
+++ b/x/distribution/types/pool.go
@@ -130,13 +130,13 @@ func (pools Pools) CoinsOf(denom string) sdk.Coins {
 
 	default:
 		midIdx := len(pools) / 2 // 2:1, 3:1, 4:2
-		coin := pools[midIdx]
+		pool := pools[midIdx]
 
 		switch {
-		case denom < coin.Denom:
+		case denom < pool.Denom:
 			return pools[:midIdx].CoinsOf(denom)
-		case denom == coin.Denom:
-			return coin.Coins
+		case denom == pool.Denom:
+			return pool.Coins
 		default:
 			return pools[midIdx+1:].CoinsOf(denom)
 		}
@@ -202,7 +202,12 @@ var _ sort.Interface = Pools{}
 
 // Sort is a helper function to sort the set of p in-place
 func (p Pools) Sort() Pools {
-	sort.Sort(p)
+	// sort.Sort does a costly runtime copy as part of `runtime.convTSlice`
+	// So we avoid this heap allocation if len(pools) <= 1. In the future, we should hopefully find
+	// a strategy to always avoid this.
+	if len(p) > 1 {
+		sort.Sort(p)
+	}
 	return p
 }
 
@@ -211,6 +216,7 @@ func (p Pools) isSorted() bool {
 		if p[i-1].Denom > p[i].Denom {
 			return false
 		}
+
 	}
 	return true
 }
@@ -220,7 +226,8 @@ func (p Pools) isSorted() bool {
 
 // NewPool return new pool instance
 func NewPool(denom string, coins sdk.Coins) Pool {
-	return Pool{denom, coins}
+	// use NewCoins to ensure the coins are sorted
+	return Pool{denom, sdk.NewCoins(coins...)}
 }
 
 // IsEmpty returns wether the pool coins are empty or not
@@ -249,7 +256,7 @@ func (pool Pool) Sub(poolB Pool) Pool {
 }
 
 func removeZeroPools(pools Pools) Pools {
-	result := make([]Pool, 0, len(pools))
+	result := make(Pools, 0, len(pools))
 
 	for _, pool := range pools {
 		if !pool.IsEmpty() {
@@ -257,13 +264,13 @@ func removeZeroPools(pools Pools) Pools {
 		}
 	}
 
-	return result
+	return result.Sort()
 }
 
 // IsEqual returns true if the two sets of Pools have the same value.
 func (pool Pool) IsEqual(other Pool) bool {
 	if pool.Denom != other.Denom {
-		panic(fmt.Sprintf("invalid pool denominations; %s, %s", pool.Denom, other.Denom))
+		return false
 	}
 
 	return pool.Coins.Equal(other.Coins)
diff --git a/x/distribution/types/pool_test.go b/x/distribution/types/pool_test.go
new file mode 100644
index 00000000..9046c3be
--- /dev/null
+++ b/x/distribution/types/pool_test.go
@@ -0,0 +1,268 @@
+package types_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+
+	"cosmossdk.io/math"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/initia-labs/initia/x/distribution/types"
+)
+
+var (
+	testPoolDenom0 = "pool0"
+	testPoolDenom1 = "pool1"
+	testPoolDenom2 = "pool2"
+	testPoolDenom4 = "pool4"
+)
+
+var (
+	testCoinDenom1 = "coin1"
+	testCoinDenom2 = "coin2"
+)
+
+type poolTestSuite struct {
+	suite.Suite
+	ca0, ca1, ca2, ca4, cm0, cm1, cm2, cm4 sdk.Coin
+	emptyCoins                             sdk.Coins
+
+	pool0, pool1, pool2, pool4                  types.Pool
+	pool000, pool111, pool222, pool444          types.Pool
+	pool101, pool110, pool122, pool211, pool244 types.Pool
+	emptyPools                                  types.Pools
+}
+
+func TestPoolTestSuite(t *testing.T) {
+	suite.Run(t, new(poolTestSuite))
+}
+
+func (s *poolTestSuite) SetupSuite() {
+	zero := math.NewInt(0)
+	one := math.OneInt()
+	two := math.NewInt(2)
+	four := math.NewInt(4)
+
+	s.ca0, s.ca1, s.ca2, s.ca4 = sdk.NewCoin(testCoinDenom1, zero), sdk.NewCoin(testCoinDenom1, one), sdk.NewCoin(testCoinDenom1, two), sdk.NewCoin(testCoinDenom1, four)
+	s.cm0, s.cm1, s.cm2, s.cm4 = sdk.NewCoin(testCoinDenom2, zero), sdk.NewCoin(testCoinDenom2, one), sdk.NewCoin(testCoinDenom2, two), sdk.NewCoin(testCoinDenom2, four)
+	s.emptyCoins = sdk.Coins{}
+
+	s.pool0 = types.NewPool(testPoolDenom0, sdk.Coins{})
+	s.pool1 = types.NewPool(testPoolDenom1, sdk.Coins{})
+	s.pool2 = types.NewPool(testPoolDenom2, sdk.Coins{})
+	s.pool4 = types.NewPool(testPoolDenom4, sdk.Coins{})
+
+	s.pool000 = types.NewPool(testPoolDenom0, sdk.NewCoins(s.ca0, s.cm0))
+	s.pool111 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca1, s.cm1))
+	s.pool222 = types.NewPool(testPoolDenom2, sdk.NewCoins(s.ca2, s.cm2))
+	s.pool444 = types.NewPool(testPoolDenom4, sdk.NewCoins(s.ca4, s.cm4))
+
+	s.pool101 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.cm1))
+	s.pool110 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca1))
+
+	s.pool122 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca2, s.cm2))
+	s.pool211 = types.NewPool(testPoolDenom2, sdk.NewCoins(s.ca1, s.cm1))
+	s.pool244 = types.NewPool(testPoolDenom2, sdk.NewCoins(s.ca4, s.cm4))
+
+	s.emptyPools = types.Pools{}
+}
+
+func (s *poolTestSuite) TestIsEqualPool() {
+	coins11 := sdk.NewCoins(sdk.NewInt64Coin(testCoinDenom1, 1), sdk.NewInt64Coin(testCoinDenom2, 1))
+	coins12 := sdk.NewCoins(sdk.NewInt64Coin(testCoinDenom1, 1), sdk.NewInt64Coin(testCoinDenom2, 2))
+
+	cases := []struct {
+		inputOne types.Pool
+		inputTwo types.Pool
+		expected bool
+	}{
+		{types.NewPool(testPoolDenom1, sdk.NewCoins(coins11...)), types.NewPool(testPoolDenom1, sdk.NewCoins(coins11...)), true},
+		{types.NewPool(testPoolDenom1, sdk.NewCoins(coins11...)), types.NewPool(testPoolDenom2, sdk.NewCoins(coins11...)), false},
+		{types.NewPool(testPoolDenom1, sdk.NewCoins(coins11...)), types.NewPool(testPoolDenom1, sdk.NewCoins(coins12...)), false},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.IsEqual(tc.inputTwo)
+		s.Require().Equal(tc.expected, res, "pool equality relation is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestIsEmptyPool() {
+	cases := []struct {
+		input    types.Pool
+		expected bool
+	}{
+		{types.NewPool(testPoolDenom1, s.emptyCoins), true},
+		{types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca1, s.cm1)), false},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.IsEmpty()
+		s.Require().Equal(tc.expected, res, "pool emptiness is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestAddPool() {
+	cases := []struct {
+		inputOne    types.Pool
+		inputTwo    types.Pool
+		expected    types.Pool
+		shouldPanic bool
+	}{
+		{s.pool111, s.pool111, s.pool122, false},
+		{s.pool101, s.pool110, s.pool111, false},
+		{types.NewPool(testPoolDenom1, sdk.Coins{}), s.pool111, s.pool111, false},
+		{s.pool111, s.pool211, s.pool111, true},
+	}
+
+	for tcIndex, tc := range cases {
+		if tc.shouldPanic {
+			s.Require().Panics(func() { tc.inputOne.Add(tc.inputTwo) })
+		} else {
+			res := tc.inputOne.Add(tc.inputTwo)
+			s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+		}
+	}
+}
+
+func (s *poolTestSuite) TestSubPool() {
+	cases := []struct {
+		inputOne    types.Pool
+		inputTwo    types.Pool
+		expected    types.Pool
+		shouldPanic bool
+	}{
+		{s.pool122, s.pool111, s.pool111, false},
+		{s.pool111, s.pool110, s.pool101, false},
+		{s.pool111, s.pool111, types.NewPool(testPoolDenom1, sdk.Coins{}), false},
+		{s.pool111, s.pool211, s.pool111, true},
+	}
+
+	for tcIndex, tc := range cases {
+		if tc.shouldPanic {
+			s.Require().Panics(func() { tc.inputOne.Sub(tc.inputTwo) })
+		} else {
+			res := tc.inputOne.Sub(tc.inputTwo)
+			s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+		}
+	}
+}
+
+func (s *poolTestSuite) TestNewPoolsSorted() {
+	cases := []struct {
+		input    types.Pools
+		expected types.Pools
+	}{
+		{types.NewPools(s.pool111, s.pool222), types.Pools{s.pool111, s.pool222}},
+		{types.NewPools(s.pool444, s.pool111), types.Pools{s.pool111, s.pool444}},
+	}
+
+	for tcIndex, tc := range cases {
+		s.Require().Equal(tc.input.IsEqual(tc.expected), true, "pools are not sorted, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestAddPools() {
+	cases := []struct {
+		inputOne types.Pools
+		inputTwo types.Pools
+		expected types.Pools
+	}{
+		{types.NewPools(s.pool111, s.pool222), types.Pools{}, types.NewPools(s.pool111, s.pool222)},
+		{types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool122, s.pool244)},
+		{types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool222, s.pool111), types.NewPools(s.pool122, s.pool244)},
+		{types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool444, s.pool000), types.NewPools(s.pool000, s.pool222, s.pool111, s.pool444)},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.Add(tc.inputTwo...)
+
+		s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestSubPools() {
+	cases := []struct {
+		inputOne types.Pools
+		inputTwo types.Pools
+		expected types.Pools
+	}{
+		{types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool111, s.pool222), types.Pools{}},
+		{types.NewPools(s.pool122, s.pool244), types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool111, s.pool222)},
+		{types.NewPools(s.pool122, s.pool244), types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool222, s.pool111)},
+		{types.NewPools(s.pool000, s.pool222, s.pool111, s.pool444), types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool444, s.pool000)},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.Sub(tc.inputTwo)
+
+		s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestIsAnyNegativePools() {
+	cases := []struct {
+		input    types.Pools
+		expected bool
+	}{
+		{types.NewPools(s.pool111, s.pool222, s.pool444), false},
+		{types.NewPools(types.Pool{"test", sdk.Coins{sdk.Coin{"testdenom", math.NewInt(-10)}}}), true},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.IsAnyNegative()
+		s.Require().Equal(tc.expected, res, "negative pool coins check is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestCoinsOfPools() {
+	pools := types.NewPools(s.pool111, s.pool222, s.pool444)
+
+	cases := []struct {
+		input    string
+		expected sdk.Coins
+	}{
+		{testPoolDenom1, sdk.NewCoins(s.ca1, s.cm1)},
+		{testPoolDenom2, sdk.NewCoins(s.ca2, s.cm2)},
+		{testPoolDenom4, sdk.NewCoins(s.ca4, s.cm4)},
+	}
+
+	for tcIndex, tc := range cases {
+		res := pools.CoinsOf(tc.input)
+		s.Require().True(tc.expected.Equal(res), "pool coins retrieval is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestIsEmptyPools() {
+	cases := []struct {
+		input    types.Pools
+		expected bool
+	}{
+		{types.NewPools(s.pool0), true},
+		{types.NewPools(s.pool111), false},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.IsEmpty()
+		s.Require().Equal(tc.expected, res, "pool emptiness is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *poolTestSuite) TestIsEqualPools() {
+	cases := []struct {
+		inputOne types.Pools
+		inputTwo types.Pools
+		expected bool
+	}{
+		{types.NewPools(s.pool000, s.pool111), types.NewPools(s.pool111), true},
+		{types.NewPools(s.pool111, s.pool222), types.NewPools(s.pool111), false},
+		{types.NewPools(s.pool111, s.pool222), types.Pools{s.pool111, s.pool222, s.pool000}, false}, // should we delete empty pool?
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.IsEqual(tc.inputTwo)
+		s.Require().Equal(tc.expected, res, "pools equality relation is incorrect, tc #%d", tcIndex)
+	}
+}

From 4f6b87bfbf852c56d8b4fda4813afd5b7f98bfc0 Mon Sep 17 00:00:00 2001
From: sh-cha <snghuncha@gmail.com>
Date: Tue, 25 Jun 2024 17:23:30 +0900
Subject: [PATCH 3/5] add dec pool test & fix something that could be logically
 wrong

---
 x/distribution/types/dec_pool.go      |  97 +++++----
 x/distribution/types/dec_pool_test.go | 280 ++++++++++++++++++++++++++
 x/distribution/types/pool.go          |   4 +-
 3 files changed, 342 insertions(+), 39 deletions(-)
 create mode 100644 x/distribution/types/dec_pool_test.go

diff --git a/x/distribution/types/dec_pool.go b/x/distribution/types/dec_pool.go
index dc54c98d..8c93ff1d 100644
--- a/x/distribution/types/dec_pool.go
+++ b/x/distribution/types/dec_pool.go
@@ -3,6 +3,7 @@ package types
 import (
 	"fmt"
 	"sort"
+	"strings"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -11,21 +12,14 @@ import (
 // rewards pools for multi-token staking
 type DecPools []DecPool
 
-// NewDecPools creates a new DecPools instance
-func NewDecPools(pools ...DecPool) DecPools {
-	return removeZeroDecPools(pools)
-}
-
 // NewDecPoolsFromPools create DecPools from Pools
 func NewDecPoolsFromPools(pools Pools) DecPools {
 	decPools := DecPools{}
 	for _, p := range pools {
-		if !p.IsEmpty() {
-			decPools = append(decPools, NewDecPool(p.Denom, sdk.NewDecCoinsFromCoins(p.Coins...)))
-		}
+		decPools = append(decPools, NewDecPool(p.Denom, sdk.NewDecCoinsFromCoins(p.Coins...)))
 	}
 
-	return decPools
+	return decPools.Sort()
 }
 
 // Sum returns sum of pool tokens
@@ -43,38 +37,59 @@ func (pools DecPools) Add(poolsB ...DecPool) DecPools {
 }
 
 // Add will perform addition of two DecPools sets.
-func (pools DecPools) safeAdd(poolsB DecPools) (coalesced DecPools) {
-	// probably the best way will be to make DecPools and interface and hide the structure
-	// definition (type alias)
+func (pools DecPools) safeAdd(poolsB DecPools) DecPools {
 	if !pools.isSorted() {
-		panic("DecPools (self) must be sorted")
+		panic("Pools (self) must be sorted")
 	}
 	if !poolsB.isSorted() {
-		panic("Wrong argument: DecPools must be sorted")
+		panic("Wrong argument: Pools must be sorted")
 	}
 
-	uniqPools := make(map[string]DecPools, len(pools)+len(poolsB))
-	// Traverse all the pools for each of the pools and poolsB.
-	for _, pL := range []DecPools{pools, poolsB} {
-		for _, p := range pL {
-			uniqPools[p.Denom] = append(uniqPools[p.Denom], p)
+	sum := ([]DecPool)(nil)
+	indexA, indexB := 0, 0
+	lenA, lenB := len(pools), len(poolsB)
+
+	for {
+		if indexA == lenA {
+			if indexB == lenB {
+				// return nil pools if both sets are empty
+				return sum
+			}
+
+			// return set B (excluding zero pools) if set A is empty
+			return append(sum, removeZeroDecPools(poolsB[indexB:])...)
+		} else if indexB == lenB {
+			// return set A (excluding zero pools) if set B is empty
+			return append(sum, removeZeroDecPools(pools[indexA:])...)
 		}
-	}
 
-	for denom, pL := range uniqPools { //#nosec
-		comboPool := DecPool{Denom: denom, DecCoins: sdk.DecCoins{}}
-		for _, p := range pL {
-			comboPool = comboPool.Add(p)
-		}
-		if !comboPool.IsEmpty() {
-			coalesced = append(coalesced, comboPool)
+		poolA, poolB := pools[indexA], poolsB[indexB]
+
+		switch strings.Compare(poolA.Denom, poolB.Denom) {
+		case -1: // pool A denom < pool B denom
+			if !poolA.IsEmpty() {
+				sum = append(sum, poolA)
+			}
+
+			indexA++
+
+		case 0: // pool A denom == pool B denom
+			res := poolA.Add(poolB)
+			if !res.IsEmpty() {
+				sum = append(sum, res)
+			}
+
+			indexA++
+			indexB++
+
+		case 1: // pool A denom > pool B denom
+			if !poolB.IsEmpty() {
+				sum = append(sum, poolB)
+			}
+
+			indexB++
 		}
 	}
-	if coalesced == nil {
-		return DecPools{}
-	}
-
-	return coalesced.Sort()
 }
 
 // Sub subtracts a set of DecPools from another (adds the inverse).
@@ -248,7 +263,12 @@ var _ sort.Interface = DecPools{}
 
 // Sort is a helper function to sort the set of p in-place
 func (p DecPools) Sort() DecPools {
-	sort.Sort(p)
+	// sort.Sort does a costly runtime copy as part of `runtime.convTSlice`
+	// So we avoid this heap allocation if len(dec pools) <= 1. In the future, we should hopefully find
+	// a strategy to always avoid this.
+	if len(p) > 1 {
+		sort.Sort(p)
+	}
 	return p
 }
 
@@ -257,6 +277,7 @@ func (p DecPools) isSorted() bool {
 		if p[i-1].Denom > p[i].Denom {
 			return false
 		}
+
 	}
 	return true
 }
@@ -266,7 +287,8 @@ func (p DecPools) isSorted() bool {
 
 // NewDecPool return new pool instance
 func NewDecPool(denom string, coins sdk.DecCoins) DecPool {
-	return DecPool{denom, coins}
+	// use NewDecCoins to ensure the coins are sorted
+	return DecPool{denom, sdk.NewDecCoins(coins...)}
 }
 
 // IsEmpty returns wether the pool coins are empty or not
@@ -298,11 +320,11 @@ func (pool DecPool) Sub(poolB DecPool) DecPool {
 // change. Note, the change may be zero.
 func (pool DecPool) TruncateDecimal() (Pool, DecPool) {
 	truncated, change := pool.DecCoins.TruncateDecimal()
-	return NewPool(pool.Denom, truncated), NewDecPool(pool.Denom, change)
+	return NewPool(pool.Denom, sdk.NewCoins(truncated...)), NewDecPool(pool.Denom, change)
 }
 
 func removeZeroDecPools(pools DecPools) DecPools {
-	result := make([]DecPool, 0, len(pools))
+	result := make(DecPools, 0, len(pools))
 
 	for _, pool := range pools {
 		if !pool.IsEmpty() {
@@ -316,7 +338,8 @@ func removeZeroDecPools(pools DecPools) DecPools {
 // IsEqual returns true if the two sets of DecPools have the same value.
 func (pool DecPool) IsEqual(other DecPool) bool {
 	if pool.Denom != other.Denom {
-		panic(fmt.Sprintf("invalid pool denominations; %s, %s", pool.Denom, other.Denom))
+		// panic(fmt.Sprintf("invalid pool denominations; %s, %s", pool.Denom, other.Denom))
+		return false
 	}
 
 	return pool.DecCoins.Equal(other.DecCoins)
diff --git a/x/distribution/types/dec_pool_test.go b/x/distribution/types/dec_pool_test.go
new file mode 100644
index 00000000..fbc78a17
--- /dev/null
+++ b/x/distribution/types/dec_pool_test.go
@@ -0,0 +1,280 @@
+package types_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+
+	"cosmossdk.io/math"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/initia-labs/initia/x/distribution/types"
+)
+
+type decpoolTestSuite struct {
+	suite.Suite
+	ca0, ca1, ca2, ca4, cm0, cm1, cm2, cm4 sdk.Coin
+	emptyCoins                             sdk.Coins
+
+	pool0, pool1, pool2, pool4                  types.Pool
+	pool000, pool111, pool222, pool444          types.Pool
+	pool101, pool110, pool122, pool211, pool244 types.Pool
+	emptyPools                                  types.Pools
+
+	decpool0, decpool1, decpool2, decpool4                     types.DecPool
+	decpool000, decpool111, decpool222, decpool444             types.DecPool
+	decpool101, decpool110, decpool122, decpool211, decpool244 types.DecPool
+	emptyDecPools                                              types.DecPools
+}
+
+func TestDecPoolTestSuite(t *testing.T) {
+	suite.Run(t, new(decpoolTestSuite))
+}
+
+func (s *decpoolTestSuite) SetupSuite() {
+	zero := math.NewInt(0)
+	one := math.OneInt()
+	two := math.NewInt(2)
+	four := math.NewInt(4)
+
+	s.ca0, s.ca1, s.ca2, s.ca4 = sdk.NewCoin(testCoinDenom1, zero), sdk.NewCoin(testCoinDenom1, one), sdk.NewCoin(testCoinDenom1, two), sdk.NewCoin(testCoinDenom1, four)
+	s.cm0, s.cm1, s.cm2, s.cm4 = sdk.NewCoin(testCoinDenom2, zero), sdk.NewCoin(testCoinDenom2, one), sdk.NewCoin(testCoinDenom2, two), sdk.NewCoin(testCoinDenom2, four)
+	s.emptyCoins = sdk.Coins{}
+
+	s.pool0 = types.NewPool(testPoolDenom0, sdk.Coins{})
+	s.pool1 = types.NewPool(testPoolDenom1, sdk.Coins{})
+	s.pool2 = types.NewPool(testPoolDenom2, sdk.Coins{})
+	s.pool4 = types.NewPool(testPoolDenom4, sdk.Coins{})
+
+	s.pool000 = types.NewPool(testPoolDenom0, sdk.NewCoins(s.ca0, s.cm0))
+	s.pool111 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca1, s.cm1))
+	s.pool222 = types.NewPool(testPoolDenom2, sdk.NewCoins(s.ca2, s.cm2))
+	s.pool444 = types.NewPool(testPoolDenom4, sdk.NewCoins(s.ca4, s.cm4))
+
+	s.pool101 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.cm1))
+	s.pool110 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca1))
+
+	s.pool122 = types.NewPool(testPoolDenom1, sdk.NewCoins(s.ca2, s.cm2))
+	s.pool211 = types.NewPool(testPoolDenom2, sdk.NewCoins(s.ca1, s.cm1))
+	s.pool244 = types.NewPool(testPoolDenom2, sdk.NewCoins(s.ca4, s.cm4))
+
+	s.emptyPools = types.Pools{}
+
+	s.decpool0 = types.NewDecPool(testPoolDenom0, sdk.DecCoins{})
+	s.decpool1 = types.NewDecPool(testPoolDenom1, sdk.DecCoins{})
+	s.decpool2 = types.NewDecPool(testPoolDenom2, sdk.DecCoins{})
+	s.decpool4 = types.NewDecPool(testPoolDenom4, sdk.DecCoins{})
+
+	s.decpool000 = types.NewDecPool(testPoolDenom0, sdk.NewDecCoinsFromCoins(s.ca0, s.cm0))
+	s.decpool111 = types.NewDecPool(testPoolDenom1, sdk.NewDecCoinsFromCoins(s.ca1, s.cm1))
+	s.decpool222 = types.NewDecPool(testPoolDenom2, sdk.NewDecCoinsFromCoins(s.ca2, s.cm2))
+	s.decpool444 = types.NewDecPool(testPoolDenom4, sdk.NewDecCoinsFromCoins(s.ca4, s.cm4))
+
+	s.decpool101 = types.NewDecPool(testPoolDenom1, sdk.NewDecCoinsFromCoins(s.cm1))
+	s.decpool110 = types.NewDecPool(testPoolDenom1, sdk.NewDecCoinsFromCoins(s.ca1))
+
+	s.decpool122 = types.NewDecPool(testPoolDenom1, sdk.NewDecCoinsFromCoins(s.ca2, s.cm2))
+	s.decpool211 = types.NewDecPool(testPoolDenom2, sdk.NewDecCoinsFromCoins(s.ca1, s.cm1))
+	s.decpool244 = types.NewDecPool(testPoolDenom2, sdk.NewDecCoinsFromCoins(s.ca4, s.cm4))
+
+	s.emptyDecPools = types.DecPools{}
+}
+
+func (s *decpoolTestSuite) TestIsEqualPool() {
+	coins11 := sdk.NewDecCoins(sdk.NewInt64DecCoin(testCoinDenom1, 1), sdk.NewInt64DecCoin(testCoinDenom2, 1))
+	coins12 := sdk.NewDecCoins(sdk.NewInt64DecCoin(testCoinDenom1, 1), sdk.NewInt64DecCoin(testCoinDenom2, 2))
+
+	cases := []struct {
+		inputOne types.DecPool
+		inputTwo types.DecPool
+		expected bool
+	}{
+		{types.NewDecPool(testPoolDenom1, sdk.NewDecCoins(coins11...)), types.NewDecPool(testPoolDenom1, sdk.NewDecCoins(coins11...)), true},
+		{types.NewDecPool(testPoolDenom1, sdk.NewDecCoins(coins11...)), types.NewDecPool(testPoolDenom2, sdk.NewDecCoins(coins11...)), false},
+		{types.NewDecPool(testPoolDenom1, sdk.NewDecCoins(coins11...)), types.NewDecPool(testPoolDenom1, sdk.NewDecCoins(coins12...)), false},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.IsEqual(tc.inputTwo)
+		s.Require().Equal(tc.expected, res, "pool equality relation is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestIsEmptyPool() {
+	cases := []struct {
+		input    types.DecPool
+		expected bool
+	}{
+		{types.NewDecPool(testPoolDenom1, sdk.DecCoins{}), true},
+		{types.NewDecPool(testPoolDenom1, sdk.NewDecCoins(sdk.NewDecCoinFromCoin(s.ca1), sdk.NewDecCoinFromCoin(s.cm1))), false},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.IsEmpty()
+		s.Require().Equal(tc.expected, res, "pool emptiness is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestAddPool() {
+	cases := []struct {
+		inputOne    types.DecPool
+		inputTwo    types.DecPool
+		expected    types.DecPool
+		shouldPanic bool
+	}{
+		{s.decpool111, s.decpool111, s.decpool122, false},
+		{s.decpool101, s.decpool110, s.decpool111, false},
+		{types.NewDecPool(testPoolDenom1, sdk.DecCoins{}), s.decpool111, s.decpool111, false},
+		{s.decpool111, s.decpool211, s.decpool111, true},
+	}
+
+	for tcIndex, tc := range cases {
+		if tc.shouldPanic {
+			s.Require().Panics(func() { tc.inputOne.Add(tc.inputTwo) })
+		} else {
+			res := tc.inputOne.Add(tc.inputTwo)
+			s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+		}
+	}
+}
+
+func (s *decpoolTestSuite) TestSubPool() {
+	cases := []struct {
+		inputOne    types.DecPool
+		inputTwo    types.DecPool
+		expected    types.DecPool
+		shouldPanic bool
+	}{
+		{s.decpool122, s.decpool111, s.decpool111, false},
+		{s.decpool111, s.decpool110, s.decpool101, false},
+		{s.decpool111, s.decpool111, types.DecPool{testPoolDenom1, sdk.DecCoins(nil)}, false},
+		{s.decpool111, s.decpool211, s.decpool111, true},
+	}
+
+	for tcIndex, tc := range cases {
+		if tc.shouldPanic {
+			s.Require().Panics(func() { tc.inputOne.Sub(tc.inputTwo) })
+		} else {
+			res := tc.inputOne.Sub(tc.inputTwo)
+			s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+		}
+	}
+}
+
+func (s *decpoolTestSuite) TestNewPoolsSorted() {
+	cases := []struct {
+		input    types.DecPools
+		expected types.DecPools
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.DecPools{s.decpool111, s.decpool222}},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool444, s.pool111)), types.DecPools{s.decpool111, s.decpool444}},
+	}
+
+	for tcIndex, tc := range cases {
+		s.Require().Equal(tc.input.IsEqual(tc.expected), true, "pools are not sorted, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestAddPools() {
+	cases := []struct {
+		inputOne types.DecPools
+		inputTwo types.DecPools
+		expected types.DecPools
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.DecPools{}, types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222))},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool122, s.pool244))},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool222, s.pool111)), types.NewDecPoolsFromPools(types.NewPools(s.pool122, s.pool244))},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool444, s.pool000)), types.NewDecPoolsFromPools(types.NewPools(s.pool000, s.pool222, s.pool111, s.pool444))},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.Add(tc.inputTwo...)
+
+		s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestSubPools() {
+	cases := []struct {
+		inputOne types.DecPools
+		inputTwo types.DecPools
+		expected types.DecPools
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.DecPools(nil)},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool122, s.pool244)), types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222))},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool122, s.pool244)), types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool222, s.pool111))},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool000, s.pool222, s.pool111, s.pool444)), types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool444, s.pool000))},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.Sub(tc.inputTwo)
+
+		s.Require().Equal(tc.expected, res, "sum of pools is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestIsAnyNegativePools() {
+	cases := []struct {
+		input    types.DecPools
+		expected bool
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222, s.pool444)), false},
+		{types.DecPools{types.DecPool{"test", sdk.DecCoins{sdk.DecCoin{"testdenom", math.LegacyNewDecFromInt(math.NewInt(-10))}}}}, true},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.IsAnyNegative()
+		s.Require().Equal(tc.expected, res, "negative pool coins check is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestCoinsOfPools() {
+	pools := types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222, s.pool444))
+
+	cases := []struct {
+		input    string
+		expected sdk.DecCoins
+	}{
+		{testPoolDenom1, sdk.NewDecCoinsFromCoins(s.ca1, s.cm1)},
+		{testPoolDenom2, sdk.NewDecCoinsFromCoins(s.ca2, s.cm2)},
+		{testPoolDenom4, sdk.NewDecCoinsFromCoins(s.ca4, s.cm4)},
+	}
+
+	for tcIndex, tc := range cases {
+		res := pools.CoinsOf(tc.input)
+		s.Require().True(tc.expected.Equal(res), "pool coins retrieval is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestIsEmptyPools() {
+	cases := []struct {
+		input    types.DecPools
+		expected bool
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool0)), true},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111)), false},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.IsEmpty()
+		s.Require().Equal(tc.expected, res, "pool emptiness is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestIsEqualPools() {
+	cases := []struct {
+		inputOne types.DecPools
+		inputTwo types.DecPools
+		expected bool
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool000, s.pool111)), types.NewDecPoolsFromPools(types.NewPools(s.pool111)), true},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.NewPools(s.pool111)), false},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)), types.NewDecPoolsFromPools(types.Pools{s.pool111, s.pool222, s.pool000}), false}, // should we delete empty pool?
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.IsEqual(tc.inputTwo)
+		s.Require().Equal(tc.expected, res, "pools equality relation is incorrect, tc #%d", tcIndex)
+	}
+}
diff --git a/x/distribution/types/pool.go b/x/distribution/types/pool.go
index 67f333c4..6a8c2143 100644
--- a/x/distribution/types/pool.go
+++ b/x/distribution/types/pool.go
@@ -13,7 +13,7 @@ type Pools []Pool
 
 // NewPools creates a new Pools instance
 func NewPools(pools ...Pool) Pools {
-	return removeZeroPools(pools)
+	return removeZeroPools(pools).Sort()
 }
 
 // Sum returns sum of pool tokens
@@ -264,7 +264,7 @@ func removeZeroPools(pools Pools) Pools {
 		}
 	}
 
-	return result.Sort()
+	return result
 }
 
 // IsEqual returns true if the two sets of Pools have the same value.

From a0b95e8f7e11dc39ef940a807bc5a7e69323fa88 Mon Sep 17 00:00:00 2001
From: sh-cha <snghuncha@gmail.com>
Date: Tue, 25 Jun 2024 17:48:31 +0900
Subject: [PATCH 4/5] add dec test

---
 x/distribution/types/dec_pool_test.go | 64 +++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/x/distribution/types/dec_pool_test.go b/x/distribution/types/dec_pool_test.go
index fbc78a17..6becaafb 100644
--- a/x/distribution/types/dec_pool_test.go
+++ b/x/distribution/types/dec_pool_test.go
@@ -278,3 +278,67 @@ func (s *decpoolTestSuite) TestIsEqualPools() {
 		s.Require().Equal(tc.expected, res, "pools equality relation is incorrect, tc #%d", tcIndex)
 	}
 }
+
+func (s *decpoolTestSuite) TestSumPools() {
+	cases := []struct {
+		input    types.DecPools
+		expected sdk.DecCoins
+	}{
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool122, s.pool222)), sdk.NewDecCoinsFromCoins(s.ca4, s.cm4)},
+		{types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool111)), sdk.NewDecCoinsFromCoins(s.ca2, s.cm2)},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.input.Sum()
+		s.Require().True(tc.expected.Equal(res), "sum of pools is incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestTruncatePools() {
+	cases := []struct {
+		input     types.DecPools
+		expected1 types.Pools
+		expected2 types.DecPools
+	}{
+		{
+			types.DecPools{types.DecPool{"test1", sdk.DecCoins{sdk.DecCoin{"testdenom1", math.LegacyNewDecFromIntWithPrec(math.NewInt(10500), 3)}}}},
+			types.Pools{types.Pool{"test1", sdk.Coins{sdk.Coin{"testdenom1", math.NewInt(10)}}}},
+			types.DecPools{types.DecPool{"test1", sdk.DecCoins{sdk.DecCoin{"testdenom1", math.LegacyNewDecFromIntWithPrec(math.NewInt(500), 3)}}}},
+		},
+		{
+			types.DecPools{types.DecPool{"test1", sdk.DecCoins{sdk.DecCoin{"testdenom1", math.LegacyNewDecFromIntWithPrec(math.NewInt(10000), 3)}}}},
+			types.Pools{types.Pool{"test1", sdk.Coins{sdk.Coin{"testdenom1", math.NewInt(10)}}}},
+			types.DecPools{},
+		},
+	}
+
+	for tcIndex, tc := range cases {
+		res1, res2 := tc.input.TruncateDecimal()
+		s.Require().True(tc.expected1.IsEqual(res1), "truncated pools are incorrect, tc #%d", tcIndex)
+		s.Require().True(tc.expected2.IsEqual(res2), "change pools are incorrect, tc #%d", tcIndex)
+	}
+}
+
+func (s *decpoolTestSuite) TestIntersectPools() {
+	cases := []struct {
+		inputOne types.DecPools
+		inputTwo types.DecPools
+		expected types.DecPools
+	}{
+		{
+			types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)),
+			types.NewDecPoolsFromPools(types.NewPools(s.pool110)),
+			types.NewDecPoolsFromPools(types.NewPools(s.pool110)),
+		},
+		{
+			types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool222)),
+			types.NewDecPoolsFromPools(types.NewPools(s.pool111, s.pool444)),
+			types.NewDecPoolsFromPools(types.NewPools(s.pool111)),
+		},
+	}
+
+	for tcIndex, tc := range cases {
+		res := tc.inputOne.Intersect(tc.inputTwo)
+		s.Require().True(tc.expected.IsEqual(res), "intersection of pools is incorrect, tc #%d", tcIndex)
+	}
+}

From ecd11b236cae18e432e745797da2b301814933a8 Mon Sep 17 00:00:00 2001
From: sh-cha <snghuncha@gmail.com>
Date: Thu, 27 Jun 2024 15:46:18 +0900
Subject: [PATCH 5/5] update dist test

---
 x/distribution/keeper/allocation_test.go |  145 ++-
 x/distribution/keeper/common_test.go     |   66 +-
 x/distribution/keeper/delegation_test.go | 1030 ++++++++++++++++------
 3 files changed, 927 insertions(+), 314 deletions(-)

diff --git a/x/distribution/keeper/allocation_test.go b/x/distribution/keeper/allocation_test.go
index 36d2d7c6..196e6e94 100644
--- a/x/distribution/keeper/allocation_test.go
+++ b/x/distribution/keeper/allocation_test.go
@@ -62,10 +62,7 @@ func TestLoadRewardWeights(t *testing.T) {
 func TestLoadBondedTokens(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar"})
-
-	// update reward weights
-	// update reward weights
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
 	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
 		{
 			Denom:  "foo",
@@ -75,14 +72,20 @@ func TestLoadBondedTokens(t *testing.T) {
 			Denom:  "bar",
 			Weight: math.LegacyNewDecWithPrec(6, 1),
 		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
 	})
 
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
 	valAddr1 := createValidatorWithCoin(ctx, input,
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000)),
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 3_000_000), sdk.NewInt64Coin("bar", 5_000_000)), 1)
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 3_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 20_000)), 1)
 	valAddr2 := createValidatorWithCoin(ctx, input,
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000)),
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 5_000_000), sdk.NewInt64Coin("bar", 3_000_000)), 2)
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 6_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
 	validator1, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
@@ -100,7 +103,7 @@ func TestLoadBondedTokens(t *testing.T) {
 	}
 	abciValB := abci.Validator{
 		Address: valConsPk2.Address(),
-		Power:   100,
+		Power:   400,
 	}
 
 	votes := []abci.VoteInfo{
@@ -123,47 +126,75 @@ func TestLoadBondedTokens(t *testing.T) {
 		if val.ValAddr == validator1.GetOperator() {
 			require.Equal(t, math.NewInt(3_000_000), val.Amount)
 		} else {
-			require.Equal(t, math.NewInt(5_000_000), val.Amount)
+			require.Equal(t, math.NewInt(6_000_000), val.Amount)
 		}
 	}
 
 	for _, val := range bondedTokens["bar"] {
-		if val.ValAddr == validator2.GetOperator() {
-			require.Equal(t, math.NewInt(3_000_000), val.Amount)
+		if val.ValAddr == validator1.GetOperator() {
+			require.Equal(t, math.NewInt(1_000_000), val.Amount)
 		} else {
-			require.Equal(t, math.NewInt(5_000_000), val.Amount)
+			require.Equal(t, math.NewInt(4_000_000), val.Amount)
 		}
 	}
-	require.Equal(t, math.NewInt(8_000_000), bondedTokensSum["foo"])
-	require.Equal(t, math.NewInt(8_000_000), bondedTokensSum["bar"])
+
+	for _, val := range bondedTokens["aaa"] {
+		if val.ValAddr == validator1.GetOperator() {
+			require.Equal(t, math.NewInt(20_000), val.Amount)
+		} else {
+			require.Equal(t, math.NewInt(10_000), val.Amount)
+		}
+	}
+	require.Equal(t, math.NewInt(9_000_000), bondedTokensSum["foo"])
+	require.Equal(t, math.NewInt(5_000_000), bondedTokensSum["bar"])
+	require.Equal(t, math.NewInt(30_000), bondedTokensSum["aaa"])
 }
 
 func TestAllocateTokensToValidatorWithCommission(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	valAddr := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar"})
+
+	// update reward weights
+	// update reward weights
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+	})
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 3_000_000), sdk.NewInt64Coin("bar", 5_000_000)), 1)
 
 	validator, err := input.StakingKeeper.Validator(ctx, valAddr)
 	require.NoError(t, err)
 
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(10)}}
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, validator, bondDenom, tokens)
-	expected := customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(5)}}}}
+	tokens := sdk.DecCoins{{Denom: "reward1", Amount: math.LegacyNewDec(10)}, {Denom: "reward2", Amount: math.LegacyNewDec(20)}}
+	input.DistKeeper.AllocateTokensToValidatorPool(ctx, validator, "bar", tokens)
+	expectedCommission := customtypes.DecPools{{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: "reward1", Amount: math.LegacyNewDec(1)}, {Denom: "reward2", Amount: math.LegacyNewDec(2)}}}}
+	expectedRewards := customtypes.DecPools{{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: "reward1", Amount: math.LegacyNewDec(9)}, {Denom: "reward2", Amount: math.LegacyNewDec(18)}}}}
 
 	// check commission
 	commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr)
 	require.NoError(t, err)
-	require.Equal(t, expected, commission.Commissions)
+	require.Equal(t, expectedCommission, commission.Commissions)
 	// check current rewards
 
 	currentRewards, err := input.DistKeeper.ValidatorCurrentRewards.Get(ctx, valAddr)
 	require.NoError(t, err)
-	require.Equal(t, expected, currentRewards.Rewards)
+	require.Equal(t, expectedRewards, currentRewards.Rewards)
 }
 
 func TestAllocateTokensToManyValidators(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
-	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar"})
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
 	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
 		{
 			Denom:  "foo",
@@ -173,14 +204,20 @@ func TestAllocateTokensToManyValidators(t *testing.T) {
 			Denom:  "bar",
 			Weight: math.LegacyNewDecWithPrec(6, 1),
 		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
 	})
 
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
 	valAddr1 := createValidatorWithCoin(ctx, input,
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000)),
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 3_000_000), sdk.NewInt64Coin("bar", 5_000_000)), 1)
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
 	valAddr2 := createValidatorWithCoin(ctx, input,
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000)),
-		sdk.NewCoins(sdk.NewInt64Coin("foo", 5_000_000), sdk.NewInt64Coin("bar", 3_000_000)), 2)
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
 	validator1, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
@@ -198,7 +235,7 @@ func TestAllocateTokensToManyValidators(t *testing.T) {
 	}
 	abciValB := abci.Validator{
 		Address: valConsPk2.Address(),
-		Power:   100,
+		Power:   400,
 	}
 
 	// assert initial state: zero outstanding rewards, zero community pool, zero commission, zero current rewards
@@ -247,23 +284,31 @@ func TestAllocateTokensToManyValidators(t *testing.T) {
 	require.NotNil(t, feeCollector)
 	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
 
-	input.DistKeeper.AllocateTokens(ctx, 200, votes)
+	input.DistKeeper.AllocateTokens(ctx, 500, votes)
 
 	// 98 outstanding rewards (100 less 2 to community pool)
 	val1OutRewards, err = input.DistKeeper.ValidatorOutstandingRewards.Get(ctx, valAddr1)
 	require.NoError(t, err)
 	require.Equal(t,
 		customtypes.DecPools{
-			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(3675, 2)}}},
-			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(147, 1)}}},
+			// 98 * (40_000 / 50_000) * (10 / 20) = 39.2
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(392, 1)}}},
+			// 98 * (1_000_000 / 5_000_000) * (6 / 20) = 5.88
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(588, 2)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) = 9.8
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 1)}}},
 		},
 		val1OutRewards.Rewards)
 	val2OutRewards, err = input.DistKeeper.ValidatorOutstandingRewards.Get(ctx, valAddr2)
 	require.NoError(t, err)
 	require.Equal(t,
 		customtypes.DecPools{
-			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(2205, 2)}}},
-			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(245, 1)}}},
+			// 98 * (10_000 / 50_000) * (10 / 20) = 9.8
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 1)}}},
+			// 98 * (4_000_000 / 5_000_000) * (6 / 20) = 23.52
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(2352, 2)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) = 9.8
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 1)}}},
 		},
 		val2OutRewards.Rewards)
 	// 2 community pool coins
@@ -271,39 +316,55 @@ func TestAllocateTokensToManyValidators(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(2)}}, feePool.CommunityPool)
 
-	// 50% commission for first proposer, (0.5 * 98%) * 100 / 2 = 24.5
+	// 10% commission for first proposer,
 	val1Commission, err = input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
 	require.Equal(t,
 		customtypes.DecPools{
-			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(18375, 3)}}},
-			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(735, 2)}}},
+			// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10) = 3.92
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(392, 2)}}},
+			// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10) = 0.588
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(588, 3)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (1 / 10) = 0.98
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 2)}}},
 		},
 		val1Commission.Commissions)
 	val2Commission, err = input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr2)
 	require.NoError(t, err)
 	require.Equal(t,
 		customtypes.DecPools{
-			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(11025, 3)}}},
-			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(1225, 2)}}},
+			// 98 * (10_000 / 50_000) * (10 / 20) * (1 / 10) = 0.98
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 2)}}},
+			// 98 * (4_000_000 / 5_000_000) * (6 / 20) * (1 / 10) = 2.352
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(2352, 3)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (1 / 10) = 0.98
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 2)}}},
 		},
 		val2Commission.Commissions)
 
-	// just staking.proportional for first proposer less commission = (0.5 * 98%) * 100 / 2 = 24.5
+	// just staking.proportional for first proposer less commission
 	val1CurRewards, err = input.DistKeeper.ValidatorCurrentRewards.Get(ctx, valAddr1)
 	require.NoError(t, err)
 	require.Equal(t,
 		customtypes.DecPools{
-			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(18375, 3)}}},
-			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(735, 2)}}},
+			// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10) = 35.28
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(3528, 2)}}},
+			// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10) = 5.292
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(5292, 3)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10) = 8.82
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(882, 2)}}},
 		},
 		val1CurRewards.Rewards)
 	val2CurRewards, err = input.DistKeeper.ValidatorCurrentRewards.Get(ctx, valAddr2)
 	require.NoError(t, err)
 	require.Equal(t,
 		customtypes.DecPools{
-			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(11025, 3)}}},
-			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(1225, 2)}}},
+			// 98 * (10_000 / 50_000) * (10 / 20) * (9 / 10) = 8.82
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(882, 2)}}},
+			// 98 * (4_000_000 / 5_000_000) * (6 / 20) * (9 / 10) = 21.168
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(21168, 3)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10) = 8.82
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(882, 2)}}},
 		},
 		val2CurRewards.Rewards)
 }
diff --git a/x/distribution/keeper/common_test.go b/x/distribution/keeper/common_test.go
index 619f51a9..3de3dc93 100644
--- a/x/distribution/keeper/common_test.go
+++ b/x/distribution/keeper/common_test.go
@@ -189,16 +189,17 @@ func (f *TestFaucet) NewFundedAccount(ctx sdk.Context, amounts ...sdk.Coin) sdk.
 }
 
 type TestKeepers struct {
-	AccountKeeper  authkeeper.AccountKeeper
-	StakingKeeper  stakingkeeper.Keeper
-	DistKeeper     distrkeeper.Keeper
-	BankKeeper     bankkeeper.Keeper
-	GovKeeper      govkeeper.Keeper
-	MoveKeeper     movekeeper.Keeper
-	DexKeeper      TestDexKeeper
-	EncodingConfig initiaappparams.EncodingConfig
-	Faucet         *TestFaucet
-	MultiStore     storetypes.CommitMultiStore
+	AccountKeeper     authkeeper.AccountKeeper
+	StakingKeeper     stakingkeeper.Keeper
+	DistKeeper        distrkeeper.Keeper
+	BankKeeper        bankkeeper.Keeper
+	GovKeeper         govkeeper.Keeper
+	MoveKeeper        movekeeper.Keeper
+	DexKeeper         TestDexKeeper
+	VotingPowerKeeper *TestVotingPowerKeeper
+	EncodingConfig    initiaappparams.EncodingConfig
+	Faucet            *TestFaucet
+	MultiStore        storetypes.CommitMultiStore
 }
 
 // createDefaultTestInput common settings for createTestInput
@@ -300,12 +301,14 @@ func _createTestInput(
 	)
 	require.NoError(t, bankKeeper.SetParams(ctx, banktypes.DefaultParams()))
 
+	votingPowerKeeper := NewTestVotingPowerKeeper()
+
 	stakingKeeper := stakingkeeper.NewKeeper(
 		appCodec,
 		runtime.NewKVStoreService(keys[stakingtypes.StoreKey]),
 		accountKeeper,
 		bankKeeper,
-		movekeeper.NewVotingPowerKeeper(moveKeeper),
+		votingPowerKeeper,
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 		vc, cc,
 	)
@@ -406,16 +409,17 @@ func _createTestInput(
 	cfg.SetAddressVerifier(initiaapp.VerifyAddressLen())
 
 	keepers := TestKeepers{
-		AccountKeeper:  accountKeeper,
-		StakingKeeper:  *stakingKeeper,
-		DistKeeper:     *distKeeper,
-		MoveKeeper:     *moveKeeper,
-		BankKeeper:     bankKeeper,
-		GovKeeper:      *govKeeper,
-		DexKeeper:      dexKeeper,
-		EncodingConfig: encodingConfig,
-		Faucet:         faucet,
-		MultiStore:     ms,
+		AccountKeeper:     accountKeeper,
+		StakingKeeper:     *stakingKeeper,
+		DistKeeper:        *distKeeper,
+		MoveKeeper:        *moveKeeper,
+		BankKeeper:        bankKeeper,
+		GovKeeper:         *govKeeper,
+		DexKeeper:         dexKeeper,
+		VotingPowerKeeper: votingPowerKeeper,
+		EncodingConfig:    encodingConfig,
+		Faucet:            faucet,
+		MultiStore:        ms,
 	}
 	return ctx, keepers
 }
@@ -512,7 +516,7 @@ func createValidatorWithCoin(
 
 // newTestMsgCreateValidator test msg creator
 func newTestMsgCreateValidator(address sdk.ValAddress, pubKey cryptotypes.PubKey, amt ...sdk.Coin) *stakingtypes.MsgCreateValidator {
-	commission := stakingtypes.NewCommissionRates(math.LegacyNewDecWithPrec(5, 1), math.LegacyNewDecWithPrec(5, 1), math.LegacyNewDec(0))
+	commission := stakingtypes.NewCommissionRates(math.LegacyNewDecWithPrec(1, 1), math.LegacyNewDecWithPrec(5, 1), math.LegacyNewDec(0))
 	msg, _ := stakingtypes.NewMsgCreateValidator(
 		address.String(), pubKey, amt,
 		stakingtypes.NewDescription("homeDir", "", "", "", ""), commission,
@@ -598,3 +602,21 @@ func (k TestDexKeeper) SwapToBase(ctx context.Context, addr sdk.AccAddress, quot
 	_, _, dummyAddr := keyPubAddr()
 	return bk.SendCoin(ctx, addr, dummyAddr, quoteCoin.Denom, quoteCoin.Amount)
 }
+
+type TestVotingPowerKeeper struct {
+	weights sdk.DecCoins
+}
+
+func NewTestVotingPowerKeeper() *TestVotingPowerKeeper {
+	return &TestVotingPowerKeeper{
+		weights: sdk.DecCoins{sdk.DecCoin{bondDenom, math.LegacyNewDec(1)}},
+	}
+}
+
+func (k *TestVotingPowerKeeper) SetVotingPowerWeights(weights sdk.DecCoins) {
+	k.weights = weights
+}
+
+func (k TestVotingPowerKeeper) GetVotingPowerWeights(_ context.Context, bondDenoms []string) (sdk.DecCoins, error) {
+	return k.weights, nil
+}
diff --git a/x/distribution/keeper/delegation_test.go b/x/distribution/keeper/delegation_test.go
index 33c49ffc..439dbb5a 100644
--- a/x/distribution/keeper/delegation_test.go
+++ b/x/distribution/keeper/delegation_test.go
@@ -5,6 +5,8 @@ import (
 
 	"cosmossdk.io/math"
 
+	abci "github.com/cometbft/cometbft/abci/types"
+	"github.com/cometbft/cometbft/proto/tendermint/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 
@@ -16,33 +18,61 @@ import (
 
 func TestCalculateRewardsBasic(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
-	valAddr1 := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
-	_, err := input.StakingKeeper.Validator(ctx, valAddr1)
-	require.NoError(t, err)
-
-	// historical count should be 2 (once for validator init, once for delegation init)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
+
+	// historical count should be 4 (once for validator init and once for delegation init per validator)
 	refCount, err := input.DistKeeper.GetValidatorHistoricalReferenceCount(ctx)
 	require.NoError(t, err)
-	require.Equal(t, uint64(2), refCount)
+	require.Equal(t, uint64(4), refCount)
 
 	// end block to bond validator and start new block
-	staking.EndBlocker(ctx, input.StakingKeeper)
+	_, err = staking.EndBlocker(ctx, input.StakingKeeper)
+	require.NoError(t, err)
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// fetch validator and delegation
 	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
 	del, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
+
 	// end period
 	endingPeriod, err := input.DistKeeper.IncrementValidatorPeriod(ctx, val)
 	require.NoError(t, err)
 
-	// historical count should be 2 still
+	// historical count should be 4 still
 	refCount, err = input.DistKeeper.GetValidatorHistoricalReferenceCount(ctx)
 	require.NoError(t, err)
-	require.Equal(t, uint64(2), refCount)
+	require.Equal(t, uint64(4), refCount)
 
 	// calculate delegation rewards
 	rewards, err := input.DistKeeper.CalculateDelegationRewards(ctx, val, del, endingPeriod)
@@ -51,10 +81,29 @@ func TestCalculateRewardsBasic(t *testing.T) {
 	// rewards should be zero
 	require.True(t, rewards.Sum().IsZero())
 
-	// allocate some rewards
-	initial := int64(10)
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial)}}
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// end period
 	endingPeriod, err = input.DistKeeper.IncrementValidatorPeriod(ctx, val)
@@ -64,36 +113,78 @@ func TestCalculateRewardsBasic(t *testing.T) {
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards should be half the tokens
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial / 2)}}}},
+		customtypes.DecPools{
+			// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10) = 35.28
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(3528, 2)}}},
+			// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10) = 5.292
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(5292, 3)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10) = 8.82
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(882, 2)}}},
+		},
 		rewards)
 
-	// commission should be the other half
 	val1Commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
+
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial / 2)}}}},
+		customtypes.DecPools{
+			// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10) = 3.92
+			{Denom: "aaa", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(392, 2)}}},
+			// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10) = 0.588
+			{Denom: "bar", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(588, 3)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (1 / 10) = 0.98
+			{Denom: "foo", DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecWithPrec(98, 2)}}},
+		},
 		val1Commission.Commissions)
 }
 
 func TestCalculateRewardsAfterSlash(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	valAddr1 := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
-	validator1, err := input.StakingKeeper.Validator(ctx, valAddr1)
-	require.NoError(t, err)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
 	// end block to bond validator and start new block
-	staking.EndBlocker(ctx, input.StakingKeeper)
+	_, err := staking.EndBlocker(ctx, input.StakingKeeper)
+	require.NoError(t, err)
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// fetch validator and delegation
 	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
 	del, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
+
 	// end period
 	endingPeriod, err := input.DistKeeper.IncrementValidatorPeriod(ctx, val)
 	require.NoError(t, err)
@@ -103,26 +194,46 @@ func TestCalculateRewardsAfterSlash(t *testing.T) {
 	// rewards should be zero
 	require.True(t, rewards.Sum().IsZero())
 
-	pubkey, err := validator1.ConsPubKey()
+	pubkey, err := val.ConsPubKey()
 	require.NoError(t, err)
 
-	// update validator for voting power update
-	_, err = input.StakingKeeper.ApplyAndReturnValidatorSetUpdates(ctx)
-	require.NoError(t, err)
-	power := validator1.GetConsensusPower(input.StakingKeeper.PowerReduction(ctx))
-	require.Equal(t, int64(1), power)
+	// start out block height
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
 	// slash the validator by 50%
-	input.StakingKeeper.Slash(ctx, pubkey.Address().Bytes(), ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	slashedTokens, err := input.StakingKeeper.Slash(ctx, pubkey.Address().Bytes(), ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
+
+	// increase block height
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
 	// retrieve validator
 	val, err = input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
 
-	// allocate some rewards
-	initial := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, 10)
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecFromInt(initial)}}
-	err = input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
 	require.NoError(t, err)
 
 	// end period
@@ -137,17 +248,49 @@ func TestCalculateRewardsAfterSlash(t *testing.T) {
 	commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
 
-	// rewards should be half the tokens
-	require.Equal(t, customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecFromInt(initial.QuoRaw(2))}}}}, rewards)
-	// commission should be the other half
-	require.Equal(t, customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecFromInt(initial.QuoRaw(2))}}}}, commission.Commissions)
+	require.Equal(t,
+		// 98 * (20_000 / 30_000) * (10 / 20) * (9 / 10)
+		// + 98 * (500_000 / 4_500_000) * (6 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 3_000_000) * (4 / 20) * (9 / 10)
+		// = 38.22
+		int64(38),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
+
+	require.Equal(t,
+		// 98 * (20_000 / 30_000) * (10 / 20) * (1 / 10)
+		// + 98 * (500_000 / 4_500_000) * (6 / 20) * (1 / 10)
+		// + 98 * (1_000_000 / 3_000_000) * (4 / 20) * (1 / 10)
+		// = 4.246666666666666
+		int64(4),
+		commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 }
 
 func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
-	valAddr1 := createValidatorWithBalance(ctx, input, 100_000_000, 100_000_000, 1)
-	_, err := input.StakingKeeper.Validator(ctx, valAddr1)
-	require.NoError(t, err)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
 	validator1, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
@@ -165,9 +308,16 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 	// fetch validator and delegation
 	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
 	del, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
+
 	// end period
 	endingPeriod, err := input.DistKeeper.IncrementValidatorPeriod(ctx, val)
 	require.NoError(t, err)
@@ -183,7 +333,9 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
 	// slash the validator by 50%
-	input.StakingKeeper.Slash(ctx, valConsAddr1, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	slashedTokens, err := input.StakingKeeper.Slash(ctx, valConsAddr1, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
 
 	// fetch the validator again
 	val, err = input.StakingKeeper.Validator(ctx, valAddr1)
@@ -192,13 +344,34 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 	// increase block height
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
-	// allocate some rewards
-	initial := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, 10)
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecFromInt(initial)}}
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// slash the validator by 50% again
-	input.StakingKeeper.Slash(ctx, valConsAddr1, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	slashedTokens, err = input.StakingKeeper.Slash(ctx, valConsAddr1, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
 
 	// fetch the validator again
 	val, err = input.StakingKeeper.Validator(ctx, valAddr1)
@@ -207,8 +380,9 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 	// increase block height
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
-	// allocate some more rewards
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// end period
 	endingPeriod, err = input.DistKeeper.IncrementValidatorPeriod(ctx, val)
@@ -218,24 +392,64 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards should be half the tokens
-	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecFromInt(initial)}}}},
-		rewards)
-
-	// commission should be the other half
 	commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
+
+	require.Equal(t,
+		// 98 * (20_000 / 30_000) * (10 / 20) * (9 / 10)
+		// + 98 * (500_000 / 4_500_000) * (6 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 3_000_000) * (4 / 20) * (9 / 10)
+		// = 38.22
+
+		// 98 * (10_000 / 20_000) * (10 / 20) * (9 / 10)
+		// + 98 * (250_000 / 4_250_000) * (6 / 20) * (9 / 10)
+		// + 98 * (500_000 / 2_500_000) * (4 / 20) * (9 / 10)
+		// = 27.134470588235295
+		int64(65),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
+
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDecFromInt(initial)}}}},
-		commission.Commissions)
+		// 98 * (20_000 / 30_000) * (10 / 20) * (1 / 10)
+		// + 98 * (500_000 / 4_500_000) * (6 / 20) * (1 / 10)
+		// + 98 * (1_000_000 / 3_000_000) * (4 / 20) * (1 / 10)
+		// = 4.246666666666666
+
+		// 98 * (10_000 / 20_000) * (10 / 20) * (1 / 10)
+		// + 98 * (250_000 / 4_250_000) * (6 / 20) * (1 / 10)
+		// + 98 * (500_000 / 2_500_000) * (4 / 20) * (1 / 10)
+		// = 3.0149411764705882
+		int64(7),
+		commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 }
 
 func TestCalculateRewardsMultiDelegator(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	// self-delegation
-	valAddr1 := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
+
 	validator, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
 	require.NoError(t, err)
 	del1, err := input.StakingKeeper.GetDelegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
@@ -244,15 +458,40 @@ func TestCalculateRewardsMultiDelegator(t *testing.T) {
 	// fetch validator and delegation
 	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
+
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
 
-	// allocate some rewards
-	initial := int64(1000)
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial)}}
-	err = input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
 	require.NoError(t, err)
 
 	// delegate to validator
-	bondCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(1_000_000)))
+	bondCoins := sdk.NewCoins(sdk.NewCoin("foo", math.NewInt(1_000_000)))
 	delAddr := input.Faucet.NewFundedAccount(ctx, bondCoins...)
 
 	shares, err := input.StakingKeeper.Delegate(ctx, delAddr, bondCoins, stakingtypes.Unbonded, validator, true)
@@ -271,7 +510,8 @@ func TestCalculateRewardsMultiDelegator(t *testing.T) {
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// allocate some more rewards
-	err = input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
 	require.NoError(t, err)
 
 	// end period
@@ -283,8 +523,18 @@ func TestCalculateRewardsMultiDelegator(t *testing.T) {
 	require.NoError(t, err)
 
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial * 3 / 4)}}}},
-		rewards)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10)
+		// + 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10)
+		// = 49.392
+
+		// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10)
+		// + 98 * (2_000_000 / 5_000_000) * (4 / 20) * (9 / 10)
+		// = 47.628
+
+		int64(97),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// calculate delegation rewards for del2
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del2, endingPeriod)
@@ -295,107 +545,195 @@ func TestCalculateRewardsMultiDelegator(t *testing.T) {
 	require.NoError(t, err)
 
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial * 1 / 4)}}}},
-		rewards)
+		// 98 * (1_000_000 / 5_000_000) * (4 / 20) * (9 / 10)
+		// = 3.528
+		int64(3),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
-	// commission should be equal to initial (50% twice)
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial)}}}},
-		commission.Commissions)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10)
+		// + 98 * (2_000_000 / 4_000_000) * (4 / 20) * (1 / 10)
+		// = 5.488
+
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10)
+		// + 98 * (3_000_000 / 5_000_000) * (4 / 20) * (1 / 10)
+		// = 5.684
+		int64(11),
+		commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 }
 
 func TestWithdrawDelegationRewardsBasic(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	// create validator with 50% commission
-	valAddr := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
-	_, err := input.StakingKeeper.GetValidator(ctx, valAddr)
-	require.NoError(t, err)
-
-	balancePower := int64(100)
-	balanceTokens := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, balancePower)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
-	// set module account coins
-	distrAcc := input.DistKeeper.GetDistributionAccount(ctx)
-	amount := sdk.NewCoins(sdk.NewCoin(bondDenom, balanceTokens))
-	err = input.BankKeeper.MintCoins(ctx, authtypes.Minter, amount)
-	require.NoError(t, err)
-	err = input.BankKeeper.SendCoinsFromModuleToModule(ctx, authtypes.Minter, distrAcc.GetName(), amount)
+	// end block to bond validator and start new block
+	_, err := staking.EndBlocker(ctx, input.StakingKeeper)
 	require.NoError(t, err)
-
-	power := int64(1)
-	valTokens := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, power)
-
-	// assert correct initial balance
-	expTokens := balanceTokens.Sub(valTokens)
-	require.Equal(t,
-		sdk.Coins{sdk.NewCoin(bondDenom, expTokens)},
-		input.BankKeeper.GetAllBalances(ctx, sdk.AccAddress(valAddr)),
-	)
-
-	// end block to bond validator
-	staking.EndBlocker(ctx, input.StakingKeeper)
-	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// fetch validator and delegation
-	val, err := input.StakingKeeper.Validator(ctx, valAddr)
+	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
+	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
+
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
 	require.NoError(t, err)
 
-	// allocate some rewards
-	initial := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, 1)
-	tokens := sdk.DecCoins{sdk.NewDecCoin(bondDenom, initial)}
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
 
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
-	// historical count should be 2 (initial + latest for delegation)
+	// historical count should be 4 (initial + latest: 2 for delegation per validator)
 	refCount, err := input.DistKeeper.GetValidatorHistoricalReferenceCount(ctx)
 	require.NoError(t, err)
-	require.Equal(t, uint64(2), refCount)
+	require.Equal(t, uint64(4), refCount)
 
 	// withdraw rewards
-	_, err = input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr), valAddr)
-	require.Nil(t, err)
+	rewards, err := input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr1), valAddr1)
+	require.NoError(t, err)
+
+	require.Equal(t,
+		customtypes.Pools{
+			// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10) = 35.28
+			{Denom: "aaa", Coins: sdk.Coins{{Denom: bondDenom, Amount: math.NewInt(35)}}},
+			// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10) = 5.292
+			{Denom: "bar", Coins: sdk.Coins{{Denom: bondDenom, Amount: math.NewInt(5)}}},
+			// 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10) = 8.82
+			{Denom: "foo", Coins: sdk.Coins{{Denom: bondDenom, Amount: math.NewInt(8)}}},
+		},
+		rewards,
+	)
 
-	// historical count should still be 2 (added one record, cleared one)
+	// historical count should still be 4 (added one record, cleared one)
 	refCount, err = input.DistKeeper.GetValidatorHistoricalReferenceCount(ctx)
 	require.NoError(t, err)
-	require.Equal(t, uint64(2), refCount)
+	require.Equal(t, uint64(4), refCount)
 
 	// assert correct balance
-	exp := balanceTokens.Sub(valTokens).Add(initial.QuoRaw(2))
 	require.Equal(t,
-		sdk.Coins{sdk.NewCoin(bondDenom, exp)},
-		input.BankKeeper.GetAllBalances(ctx, sdk.AccAddress(valAddr)),
+		rewards.Sum().AmountOf(bondDenom).Int64(),
+		input.BankKeeper.GetBalance(ctx, sdk.AccAddress(valAddr1), bondDenom).Amount.Int64(),
 	)
 
 	// withdraw commission
-	_, err = input.DistKeeper.WithdrawValidatorCommission(ctx, valAddr)
+	_, err = input.DistKeeper.WithdrawValidatorCommission(ctx, valAddr1)
 	require.Nil(t, err)
 }
 
 func TestWithdrawDelegationZeroRewards(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	// create validator with 50% commission
-	valAddr := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
-	_, err := input.StakingKeeper.GetValidator(ctx, valAddr)
-	require.NoError(t, err)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
-	balancePower := int64(1000)
-	balanceTokens := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, balancePower)
+	// fetch validator and delegation
+	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
+	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
 
-	// set module account coins
-	distrAcc := input.DistKeeper.GetDistributionAccount(ctx)
-	amount := sdk.NewCoins(sdk.NewCoin(bondDenom, balanceTokens))
-	err = input.BankKeeper.MintCoins(ctx, authtypes.Minter, amount)
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
 	require.NoError(t, err)
-	err = input.BankKeeper.SendCoinsFromModuleToModule(ctx, authtypes.Minter, distrAcc.GetName(), amount)
+
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
 	require.NoError(t, err)
-	input.AccountKeeper.SetModuleAccount(ctx, distrAcc)
 
 	// withdraw rewards -- should be 0
-	pool, err := input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr), valAddr)
+	pool, err := input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 	require.True(t, pool.Sum().IsZero(), "expected withdraw rewards to be zero")
 	require.True(t, pool.Sum().IsValid(), "expected returned coins to be valid")
@@ -404,19 +742,48 @@ func TestWithdrawDelegationZeroRewards(t *testing.T) {
 func TestCalculateRewardsAfterManySlashesInSameBlock(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	valAddr := createValidatorWithBalance(ctx, input, 100_000_000, 10_000_000, 1)
-	validator, err := input.StakingKeeper.GetValidator(ctx, valAddr)
-	require.NoError(t, err)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
-	// end block to bond validator
-	staking.EndBlocker(ctx, input.StakingKeeper)
-	// next block
+	// end block to bond validator and start new block
+	_, err := staking.EndBlocker(ctx, input.StakingKeeper)
+	require.NoError(t, err)
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// fetch validator and delegation
-	val, err := input.StakingKeeper.Validator(ctx, valAddr)
+	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
+	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
+
+	valConsPk1, err := val.ConsPubKey()
 	require.NoError(t, err)
-	del, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr), valAddr)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
+
+	del, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
 	// end period
@@ -428,32 +795,57 @@ func TestCalculateRewardsAfterManySlashesInSameBlock(t *testing.T) {
 
 	// rewards should be zero
 	require.True(t, rewards.Sum().IsZero())
-	// start out block height
-	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
-
-	// allocate some rewards
-	initial := math.LegacyNewDecFromInt(input.StakingKeeper.VotingPowerFromConsensusPower(ctx, 1))
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: initial}}
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
 
-	pubkey, err := validator.ConsPubKey()
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
+
+	pubkey, err := val.ConsPubKey()
 	require.NoError(t, err)
 	valConsAddr := pubkey.Address().Bytes()
 
+	// start out block height
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
+
 	// slash the validator by 50%
-	input.StakingKeeper.Slash(ctx, valConsAddr, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	slashedTokens, err := input.StakingKeeper.Slash(ctx, valConsAddr, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
 
 	// slash the validator by 50% again
-	input.StakingKeeper.Slash(ctx, valConsAddr, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
-
-	// fetch the validator again
-	val, err = input.StakingKeeper.Validator(ctx, valAddr)
+	slashedTokens, err = input.StakingKeeper.Slash(ctx, valConsAddr, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
 	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
+
 	// increase block height
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
-	// allocate some more rewards
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
+
+	// fetch the validator again
+	val, err = input.StakingKeeper.Validator(ctx, valAddr1)
+	require.NoError(t, err)
 
 	// end period
 	endingPeriod, err = input.DistKeeper.IncrementValidatorPeriod(ctx, val)
@@ -463,82 +855,146 @@ func TestCalculateRewardsAfterManySlashesInSameBlock(t *testing.T) {
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards should be half the tokens
-	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: initial}}}},
-		rewards)
-
 	// load commission
-	commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr)
+	commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
 
-	// commission should be the other half
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: initial}}}},
-		commission.Commissions)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10)
+		// + 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10)
+		// = 49.392
+
+		// 98 * (10_000 / 20_000) * (10 / 20) * (9 / 10)
+		// + 98 * (250_000 / 4_250_000) * (6 / 20) * (9 / 10)
+		// + 98 * (500_000 / 2_500_000) * (4 / 20) * (9 / 10)
+		// = 27.134470588235295
+		int64(76),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
+
+	require.Equal(t,
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10)
+		// + 98 * (2_000_000 / 4_000_000) * (4 / 20) * (1 / 10)
+		// = 5.488
+
+		// 98 * (10_000 / 20_000) * (10 / 20) * (1 / 10)
+		// + 98 * (250_000 / 4_250_000) * (6 / 20) * (1 / 10)
+		// + 98 * (500_000 / 2_500_000) * (4 / 20) * (1 / 10)
+		// = 3.0149411764705882
+		int64(8),
+		commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 }
 
 func TestCalculateRewardsMultiDelegatorMultiSlash(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	// self delegation
-	valAddr1 := createValidatorWithBalance(ctx, input, 100_000_000, 10_000_000, 1)
-	_, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
-	require.NoError(t, err)
-
-	valAddr2 := createValidatorWithBalance(ctx, input, 100_000_000, 1, 2)
-	_, err = input.StakingKeeper.GetValidator(ctx, valAddr2)
-	require.NoError(t, err)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
-	validator1, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
-	require.NoError(t, err)
-
-	pubkey, err := validator1.ConsPubKey()
+	// end block to bond validator and start new block
+	_, err := staking.EndBlocker(ctx, input.StakingKeeper)
 	require.NoError(t, err)
-	valConsAddr1 := pubkey.Address().Bytes()
-
-	// end block to bond validator
-	staking.EndBlocker(ctx, input.StakingKeeper)
-
-	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// fetch validator and delegation
 	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
 	del1, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
-	// allocate some rewards
-	initial := math.LegacyNewDecFromInt(input.StakingKeeper.VotingPowerFromConsensusPower(ctx, 10))
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: initial}}
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
+
+	valConsAddr := valConsPk1.Address().Bytes()
+
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// slash the validator
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
-	input.StakingKeeper.Slash(ctx, valConsAddr1, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	slashedTokens, err := input.StakingKeeper.Slash(ctx, valConsAddr, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
 	// second delegation
-	bondCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(10_000_000)))
-	shares, err := input.StakingKeeper.Delegate(ctx, sdk.AccAddress(valAddr2), bondCoins, stakingtypes.Unbonded, validator1, true)
+	bondCoins := sdk.NewCoins(sdk.NewCoin("foo", math.NewInt(1_000_000)))
+	delAddr := input.Faucet.NewFundedAccount(ctx, bondCoins...)
+
+	validator, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
 	require.NoError(t, err)
-	require.Equal(t, sdk.NewDecCoinsFromCoins(bondCoins...), shares)
-	del2, err := input.StakingKeeper.GetDelegation(ctx, sdk.AccAddress(valAddr2), valAddr1)
+	shares, err := input.StakingKeeper.Delegate(ctx, delAddr, bondCoins, stakingtypes.Unbonded, validator, true)
+	require.NoError(t, err)
+
+	// existing 2_000_000 shares for 2_000_000 / 2 tokens => new shares == 2_000_000
+	require.Equal(t, sdk.NewDecCoinsFromCoins(bondCoins...).MulDec(math.LegacyNewDec(2)), shares)
+	del2, err := input.StakingKeeper.GetDelegation(ctx, delAddr, valAddr1)
 	require.NoError(t, err)
 
 	// end block
-	staking.EndBlocker(ctx, input.StakingKeeper)
+	_, err = staking.EndBlocker(ctx, input.StakingKeeper)
+	require.NoError(t, err)
 
 	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// allocate some more rewards
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// slash the validator again
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
-	input.StakingKeeper.Slash(ctx, valConsAddr1, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	slashedTokens, err = input.StakingKeeper.Slash(ctx, valConsAddr, ctx.BlockHeight(), math.LegacyNewDecWithPrec(5, 1))
+	require.NoError(t, err)
+	require.True(t, slashedTokens.IsAllPositive(), "expected positive slashed tokens, got: %s", slashedTokens)
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 3)
 
 	// fetch updated validator
@@ -553,79 +1009,123 @@ func TestCalculateRewardsMultiDelegatorMultiSlash(t *testing.T) {
 	rewards, err := input.DistKeeper.CalculateDelegationRewards(ctx, val, del1, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards for del1 should be 5/8 initial (half initial first period, 1/8 initial second period)
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: initial.QuoInt64(2).Add(initial.QuoInt64(8))}}}},
-		rewards)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10)
+		// + 98 * (2_000_000 / 4_000_000) * (4 / 20) * (9 / 10)
+		// = 49.392
+
+		// 98 * (20_000 / 30_000) * (10 / 20) * (9 / 10)
+		// + 98 * (500_000 / 4_500_000) * (6 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 4_000_000) * (4 / 20) * (9 / 10)
+		// = 36.75
+		int64(86),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// calculate delegation rewards for del2
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del2, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards for del2 should be initial / 8
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: initial.QuoInt64(4)}}}},
-		rewards)
+		// 98 * (1_000_000 / 4_000_000) * (4 / 20) * (9 / 10)
+		// = 4.41
+		int64(4),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// load commission
 	commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
 
-	// commission should be equal to initial (twice 50% commission, unaffected by slashing)
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: initial}}}},
-		commission.Commissions)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10)
+		// + 98 * (2_000_000 / 4_000_000) * (4 / 20) * (1 / 10)
+		// = 5.488
+
+		// 98 * (20_000 / 30_000) * (10 / 20) * (1 / 10)
+		// + 98 * (500_000 / 4_500_000) * (6 / 20) * (1 / 10)
+		// + 98 * (3_000_000 / 4_000_000) * (4 / 20) * (1 / 10)
+		// = 5.06333333
+		int64(10),
+		commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 }
 
 func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	ctx, input := createDefaultTestInput(t)
 
-	// self delegation
-	valAddr1 := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 1)
-	_, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
-	require.NoError(t, err)
+	input.StakingKeeper.SetBondDenoms(ctx, []string{"foo", "bar", "aaa"})
+	setRewardWeights(t, ctx, input, []customtypes.RewardWeight{
+		{
+			Denom:  "foo",
+			Weight: math.LegacyNewDecWithPrec(4, 1),
+		},
+		{
+			Denom:  "bar",
+			Weight: math.LegacyNewDecWithPrec(6, 1),
+		},
+		{
+			Denom:  "aaa",
+			Weight: math.LegacyNewDecWithPrec(1, 0),
+		},
+	})
+
+	input.VotingPowerKeeper.SetVotingPowerWeights(sdk.NewDecCoins(sdk.NewDecCoin("foo", math.NewInt(1)), sdk.NewDecCoin("bar", math.NewInt(4)), sdk.NewDecCoin("aaa", math.NewInt(10))))
+
+	valAddr1 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 1_000_000), sdk.NewInt64Coin("aaa", 40_000)), 1)
+	valAddr2 := createValidatorWithCoin(ctx, input,
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 100_000_000), sdk.NewInt64Coin("bar", 100_000_000), sdk.NewInt64Coin("aaa", 100_000_000)),
+		sdk.NewCoins(sdk.NewInt64Coin("foo", 2_000_000), sdk.NewInt64Coin("bar", 4_000_000), sdk.NewInt64Coin("aaa", 10_000)), 2)
 
-	valAddr2 := createValidatorWithBalance(ctx, input, 100_000_000, 1_000_000, 2)
-	_, err = input.StakingKeeper.GetValidator(ctx, valAddr2)
-	require.NoError(t, err)
-
-	validator1, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
-	require.NoError(t, err)
-
-	// set module account coins
-	balancePower := int64(100)
-	balanceTokens := input.StakingKeeper.VotingPowerFromConsensusPower(ctx, balancePower)
-	distrAcc := input.DistKeeper.GetDistributionAccount(ctx)
-	amount := sdk.NewCoins(sdk.NewCoin(bondDenom, balanceTokens))
-	err = input.BankKeeper.MintCoins(ctx, authtypes.Minter, amount)
-	require.NoError(t, err)
-	err = input.BankKeeper.SendCoinsFromModuleToModule(ctx, authtypes.Minter, distrAcc.GetName(), amount)
+	// end block to bond validator and start new block
+	_, err := staking.EndBlocker(ctx, input.StakingKeeper)
 	require.NoError(t, err)
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// fetch validator and delegation
 	val, err := input.StakingKeeper.Validator(ctx, valAddr1)
 	require.NoError(t, err)
+	val2, err := input.StakingKeeper.Validator(ctx, valAddr2)
+	require.NoError(t, err)
 	del1, err := input.StakingKeeper.Delegation(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
-	// end block
-	staking.EndBlocker(ctx, input.StakingKeeper)
-	// next block
-	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
+	valConsPk1, err := val.ConsPubKey()
+	require.NoError(t, err)
+	valConsPk2, err := val2.ConsPubKey()
+	require.NoError(t, err)
 
-	// allocate some rewards (1)
-	initial := int64(20)
-	tokens := sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial)}}
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	fees := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(100)))
+	feeCollector := input.AccountKeeper.GetModuleAccount(ctx, authtypes.FeeCollectorName)
+	require.NotNil(t, feeCollector)
+	votes := []abci.VoteInfo{
+		{
+			Validator: abci.Validator{
+				Address: valConsPk1.Address(),
+				Power:   100,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+		{
+			Validator: abci.Validator{
+				Address: valConsPk2.Address(),
+				Power:   400,
+			},
+			BlockIdFlag: types.BlockIDFlagCommit,
+		},
+	}
 
 	// second delegation
-	bondCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, math.NewInt(1_000_000)))
-	shares, err := input.StakingKeeper.Delegate(ctx, sdk.AccAddress(valAddr2), bondCoins, stakingtypes.Unbonded, validator1, true)
+	bondCoins := sdk.NewCoins(sdk.NewCoin("foo", math.NewInt(1_000_000)))
+	delAddr := input.Faucet.NewFundedAccount(ctx, bondCoins...)
+
+	validator, err := input.StakingKeeper.GetValidator(ctx, valAddr1)
+	require.NoError(t, err)
+	_, err = input.StakingKeeper.Delegate(ctx, delAddr, bondCoins, stakingtypes.Unbonded, validator, true)
 	require.NoError(t, err)
-	require.Equal(t, sdk.NewDecCoinsFromCoins(bondCoins...), shares)
 
-	// fetch updated validator
-	del2, err := input.StakingKeeper.GetDelegation(ctx, sdk.AccAddress(valAddr2), valAddr1)
+	del2, err := input.StakingKeeper.GetDelegation(ctx, delAddr, valAddr1)
 	require.NoError(t, err)
 
 	// end block
@@ -635,15 +1135,17 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
-	// allocate some more rewards (2)
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	// allocate some more rewards
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// first delegator withdraws
 	_, err = input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr1), valAddr1)
 	require.NoError(t, err)
 
 	// second delegator withdraws
-	_, err = input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr2), valAddr1)
+	_, err = input.DistKeeper.WithdrawDelegationRewards(ctx, delAddr, valAddr1)
 	require.NoError(t, err)
 
 	// validator withdraws commission (1)
@@ -668,16 +1170,23 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	// rewards for del2 should be zero
 	require.True(t, rewards.Sum().IsZero())
 
-	// commission should be zero
 	val1Commission, err := input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
-	require.True(t, val1Commission.Commissions.Sum().IsZero())
+	require.Equal(t,
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10) = 3.92
+		// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10) = 0.588
+		// 98 * (3_000_000 / 5_000_000) * (4 / 20) * (1 / 10) = 1.176
+
+		// 0.92 + 0.588 + 0.176 = 1.684
+		int64(1),
+		val1Commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
-	// allocate some more rewards (3)
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// first delegator withdraws again
 	_, err = input.DistKeeper.WithdrawDelegationRewards(ctx, sdk.AccAddress(valAddr1), valAddr1)
@@ -691,23 +1200,29 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del2, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards for del2 should be 1/4 initial
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial / 4)}}}},
-		rewards)
+		// 98 * (1_000_000 / 5_000_000) * (4 / 20) * (9 / 10)
+		// = 3.5280000000000005
+		int64(3),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
-	// commission should be half initial
 	val1Commission, err = input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial / 2)}}}},
-		val1Commission.Commissions)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10) + 0.92
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10) + 0.588
+		// + 98 * (3_000_000 / 5_000_000) * (4 / 20) * (1 / 10) + 0.176
+		// 7.368
+
+		int64(7),
+		val1Commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
-	// allocate some more rewards (4)
-	input.DistKeeper.AllocateTokensToValidatorPool(ctx, val, bondDenom, tokens)
+	input.Faucet.Fund(ctx, authtypes.NewModuleAddress(authtypes.FeeCollectorName), fees...)
+	err = input.DistKeeper.AllocateTokens(ctx, 500, votes)
+	require.NoError(t, err)
 
 	// withdraw commission (2)
 	_, err = input.DistKeeper.WithdrawValidatorCommission(ctx, valAddr1)
@@ -728,22 +1243,37 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del1, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards for del1 should be 1/4 initial
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial / 4)}}}},
-		rewards)
+		// 98 * (40_000 / 50_000) * (10 / 20) * (9 / 10)
+		// + 98 * (1_000_000 / 5_000_000) * (6 / 20) * (9 / 10)
+		// + 98 * (2_000_000 / 5_000_000) * (4 / 20) * (9 / 10)
+		// = 47.628
+		int64(47),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// calculate delegation rewards for del2
 	rewards, err = input.DistKeeper.CalculateDelegationRewards(ctx, val, del2, endingPeriod)
 	require.NoError(t, err)
 
-	// rewards for del2 should be 1/2 initial
 	require.Equal(t,
-		customtypes.DecPools{{Denom: bondDenom, DecCoins: sdk.DecCoins{{Denom: bondDenom, Amount: math.LegacyNewDec(initial / 2)}}}},
-		rewards)
+		// 98 * (1_000_000 / 5_000_000) * (4 / 20) * (9 / 10)
+		// = 3.5280000000000005
+
+		// 98 * (1_000_000 / 5_000_000) * (4 / 20) * (9 / 10)
+		// = 3.5280000000000005
+
+		int64(7),
+		rewards.Sum().AmountOf(bondDenom).TruncateInt64())
 
 	// commission should be zero
 	val1Commission, err = input.DistKeeper.ValidatorAccumulatedCommissions.Get(ctx, valAddr1)
 	require.NoError(t, err)
-	require.True(t, val1Commission.Commissions.Sum().IsZero())
+	require.Equal(t,
+		// 98 * (40_000 / 50_000) * (10 / 20) * (1 / 10) * 2 + 0.92 = 8.760000000000002
+		// 98 * (1_000_000 / 5_000_000) * (6 / 20) * (1 / 10) * 2 + 0.588 = 1.7639999999999998
+		// 98 * (3_000_000 / 5_000_000) * (4 / 20) * (1 / 10) * 2 + 0.176 = 2.528
+
+		// 0.760000000000002 + 0.7639999999999998 + 0.528 = 2.052
+		int64(2),
+		val1Commission.Commissions.Sum().AmountOf(bondDenom).TruncateInt64())
 }