Add fix for prop 597 code issue (#6279)

* Add fix for prop 572 code issue

* update changelog

* Update CHANGELOG.md

* basic upgrade unit test

---------

Co-authored-by: devbot-wizard <141283918+devbot-wizard@users.noreply.github.com>
Co-authored-by: Roman <roman@osmosis.team>

CHANGELOG.md

@@ -52,6 +52,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Minor improvements & Bug Fixes
 
 ### Security
+### Misc Improvements
+
+* [#6279](https://github.com/osmosis-labs/osmosis/pull/6279) fix prop-597 introduced issue
 
 ### API Breaks
 

app/upgrades/v19/constants.go

@@ -2,6 +2,7 @@ package v19
 
 import (
 	"github.com/osmosis-labs/osmosis/v19/app/upgrades"
+	"github.com/osmosis-labs/osmosis/v19/x/gamm/types/migration"
 
 	store "github.com/cosmos/cosmos-sdk/store/types"
 )
@@ -17,3 +18,40 @@ var Upgrade = upgrades.Upgrade{
 		Deleted: []string{},
 	},
 }
+
+var Records = []migration.BalancerToConcentratedPoolLink{
+	// created at v16
+	{BalancerPoolId: 674, ClPoolId: 1066}, // DAI
+	// created at https://www.mintscan.io/osmosis/proposals/571
+	{BalancerPoolId: 837, ClPoolId: 1088},  // IST
+	{BalancerPoolId: 857, ClPoolId: 1089},  // CMST
+	{BalancerPoolId: 712, ClPoolId: 1090},  // WBTC
+	{BalancerPoolId: 773, ClPoolId: 1091},  // DOT
+	{BalancerPoolId: 9, ClPoolId: 1092},    // CRO
+	{BalancerPoolId: 3, ClPoolId: 1093},    // AKT
+	{BalancerPoolId: 812, ClPoolId: 1094},  // AXL
+	{BalancerPoolId: 584, ClPoolId: 1095},  // SCRT
+	{BalancerPoolId: 604, ClPoolId: 1096},  // STARS
+	{BalancerPoolId: 497, ClPoolId: 1097},  // JUNO
+	{BalancerPoolId: 806, ClPoolId: 1098},  // STRD
+	{BalancerPoolId: 907, ClPoolId: 1099},  // MARS
+	{BalancerPoolId: 1013, ClPoolId: 1100}, // ION
+	{BalancerPoolId: 15, ClPoolId: 1101},   // XPRT
+	{BalancerPoolId: 586, ClPoolId: 1102},  // MED
+	{BalancerPoolId: 627, ClPoolId: 1103},  // SOMM
+	{BalancerPoolId: 795, ClPoolId: 1104},  // BLD
+	{BalancerPoolId: 730, ClPoolId: 1105},  // KAVA
+	{BalancerPoolId: 7, ClPoolId: 1106},    // IRIS
+	{BalancerPoolId: 1039, ClPoolId: 1107}, // stIBCX
+	{BalancerPoolId: 5, ClPoolId: 1108},    // DVPN
+	{BalancerPoolId: 573, ClPoolId: 1109},  // BTSG
+	{BalancerPoolId: 641, ClPoolId: 1110},  // UMEE
+	{BalancerPoolId: 605, ClPoolId: 1111},  // HUAHUA
+	{BalancerPoolId: 971, ClPoolId: 1112},  // NCT
+	{BalancerPoolId: 625, ClPoolId: 1113},  // GRAV
+	// created at https://www.mintscan.io/osmosis/proposals/597
+	{BalancerPoolId: 678, ClPoolId: 1133}, // USDC
+	{BalancerPoolId: 704, ClPoolId: 1134}, // WETH
+	{BalancerPoolId: 1, ClPoolId: 1135},   // ATOM
+	{BalancerPoolId: 803, ClPoolId: 1136}, // stATOM
+}

app/upgrades/v19/upgrades.go

@@ -46,6 +46,10 @@ func CreateUpgradeHandler(
 		defaultPoolManagerParams.TakerFeeParams.DefaultTakerFee = osmomath.ZeroDec()
 		keepers.PoolManagerKeeper.SetParams(ctx, defaultPoolManagerParams)
 
+		err = keepers.GAMMKeeper.UpdateMigrationRecords(ctx, Records)
+		if err != nil {
+			return nil, err
+		}
 		return migrations, nil
 	}
 }

app/upgrades/v19/upgrades_test.go

@@ -15,7 +15,11 @@ import (
 
 	"github.com/osmosis-labs/osmosis/osmomath"
 	"github.com/osmosis-labs/osmosis/v19/app/apptesting"
+	v19 "github.com/osmosis-labs/osmosis/v19/app/upgrades/v19"
+	"github.com/osmosis-labs/osmosis/v19/x/gamm/pool-models/balancer"
 	gammtypes "github.com/osmosis-labs/osmosis/v19/x/gamm/types"
+	"github.com/osmosis-labs/osmosis/v19/x/gamm/types/migration"
+	poolincentivestypes "github.com/osmosis-labs/osmosis/v19/x/pool-incentives/types"
 
 	"github.com/osmosis-labs/osmosis/v19/x/superfluid/types"
 	superfluidtypes "github.com/osmosis-labs/osmosis/v19/x/superfluid/types"
@@ -46,9 +50,102 @@ func (s *UpgradeTestSuite) TestUpgrade() {
 	initialTokenBonded := osmomath.NewInt(100)
 	s.Setup()
 
-	// prepare superfluid delegation
-	superfluidVal, lockDenom := s.setupSuperfluidDelegation()
-	delegationBeforeV18Upgrade, _ := s.App.StakingKeeper.GetValidator(s.Ctx, superfluidVal)
+	var (
+		superfluidVal              sdk.ValAddress
+		lockDenom                  string
+		delegationBeforeV18Upgrade stakingtypes.Validator
+	)
+
+	// Get max pool ID
+	highestBalancerPoolId := uint64(0)
+	highestCLPoolId := uint64(0)
+	for _, record := range v19.Records {
+		if record.BalancerPoolId > highestBalancerPoolId {
+			highestBalancerPoolId = record.BalancerPoolId
+		}
+		if record.ClPoolId > highestCLPoolId {
+			highestCLPoolId = record.ClPoolId
+		}
+	}
+
+	// Create balancer pools
+	for i := uint64(1); i <= highestBalancerPoolId; i++ {
+
+		// 2 is a pool ID that does not exist in the linked balancer <> CL pool record.
+		// The reason is that we must have a bond denom for superfluid delegation
+		// but all balancer and CL pools to have the same ETH/USDC denoms.
+		if i == 2 {
+			// prepare superfluid delegation
+			superfluidVal, lockDenom = s.setupSuperfluidDelegation()
+			delegationBeforeV18Upgrade, _ = s.App.StakingKeeper.GetValidator(s.Ctx, superfluidVal)
+		} else {
+			s.PrepareCustomBalancerPool(
+				[]balancer.PoolAsset{
+					{
+						Token:  sdk.NewCoin(apptesting.ETH, sdk.NewInt(100000000000)),
+						Weight: sdk.NewInt(5),
+					},
+					{
+						Token:  sdk.NewCoin(apptesting.USDC, sdk.NewInt(100000000000)),
+						Weight: sdk.NewInt(5),
+					},
+				},
+				balancer.PoolParams{
+					SwapFee: sdk.ZeroDec(),
+					ExitFee: sdk.ZeroDec(),
+				},
+			)
+		}
+
+	}
+
+	// Create CL pools
+	for i := uint64(highestBalancerPoolId + 1); i <= highestCLPoolId; i++ {
+		s.PrepareConcentratedPool()
+	}
+
+	// Setup migration recods per mainnet state
+	err := s.App.GAMMKeeper.ReplaceMigrationRecords(s.Ctx, []migration.BalancerToConcentratedPoolLink{
+		{
+			BalancerPoolId: 803,
+			ClPoolId:       1136,
+		},
+	})
+	s.Require().NoError(err)
+
+	firstRecord := v19.Records[0]
+	secondRecord := v19.Records[1]
+
+	// Get CFMM gauges
+	cfmmGauges, err := s.App.PoolIncentivesKeeper.GetGaugesForCFMMPool(s.Ctx, firstRecord.BalancerPoolId)
+	s.Require().NoError(err)
+
+	// Get longest gauge duration from CFMM pool.
+	balancerLongestDurationGauge := cfmmGauges[0]
+	for i := 1; i < len(cfmmGauges); i++ {
+		if cfmmGauges[i].DistributeTo.Duration > balancerLongestDurationGauge.DistributeTo.Duration {
+			balancerLongestDurationGauge = cfmmGauges[i]
+		}
+	}
+
+	incentivesEpochDuration := s.App.IncentivesKeeper.GetEpochInfo(s.Ctx).Duration
+	clGaugeIDSecondLink, err := s.App.PoolIncentivesKeeper.GetPoolGaugeId(s.Ctx, secondRecord.ClPoolId, incentivesEpochDuration)
+	s.Require().NoError(err)
+
+	// Setup distr records
+	s.App.PoolIncentivesKeeper.SetDistrInfo(s.Ctx, poolincentivestypes.DistrInfo{
+		TotalWeight: sdk.NewInt(1000),
+		Records: []poolincentivestypes.DistrRecord{
+			{
+				GaugeId: balancerLongestDurationGauge.Id, // As a sanity check set distr record to balancer gauge id
+				Weight:  sdk.NewInt(500),
+			},
+			{
+				GaugeId: clGaugeIDSecondLink, // And to CL gauge ID, to make sure that there are no panics
+				Weight:  sdk.NewInt(500),
+			},
+		},
+	})
 
 	// run an epoch
 	s.Ctx = s.Ctx.WithBlockTime(s.Ctx.BlockTime().Add(time.Hour * 24))
@@ -65,7 +162,7 @@ func (s *UpgradeTestSuite) TestUpgrade() {
 	// broken states (current status):
 	// synth lock accumulator is set to 0
 	totalSynthLocked := s.App.SuperfluidKeeper.GetTotalSyntheticAssetsLocked(s.Ctx, stakingSyntheticDenom(lockDenom, superfluidVal.String()))
-	s.Require().True(totalSynthLocked.Equal(osmomath.ZeroInt()))
+	s.Require().Equal(totalSynthLocked.String(), osmomath.ZeroInt().String())
 
 	// superfluid delegated tokens have been undelegated from validator,
 	// only have the initial bonded amount present
@@ -86,6 +183,21 @@ func (s *UpgradeTestSuite) TestUpgrade() {
 	// also check that we have the correct superfluid staked delegation back
 	delegationAfterV19Upgrade, _ := s.App.StakingKeeper.GetValidator(s.Ctx, superfluidVal)
 	s.Require().True(delegationBeforeV18Upgrade.Tokens.Equal(delegationAfterV19Upgrade.Tokens))
+
+	// Validate that all migration records were set
+	newMigrationRecords, err := s.App.GAMMKeeper.GetAllMigrationInfo(s.Ctx)
+	s.Require().NoError(err)
+	s.Require().Len(newMigrationRecords.BalancerToConcentratedPoolLinks, len(v19.Records))
+
+	// Check that distr records point to CL gauges
+	distrInfo := s.App.PoolIncentivesKeeper.GetDistrInfo(s.Ctx)
+	s.Require().Len(distrInfo.Records, 2)
+
+	clGaugeIDFirstLink, err := s.App.PoolIncentivesKeeper.GetPoolGaugeId(s.Ctx, firstRecord.ClPoolId, incentivesEpochDuration)
+	s.Require().NoError(err)
+
+	s.Require().Equal(clGaugeIDFirstLink, distrInfo.Records[0].GaugeId)
+	s.Require().Equal(clGaugeIDSecondLink, distrInfo.Records[1].GaugeId)
 }
 
 func (s *UpgradeTestSuite) setupNormalDelegation() sdk.ValAddress {