fix: cosmwasm pool model code ID migration

CHANGELOG.md

@@ -51,6 +51,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#7509](https://github.com/osmosis-labs/osmosis/pull/7509) Distributing ProtoRev profits to the community pool and burn address
 * [#7550](https://github.com/osmosis-labs/osmosis/pull/7550) Speedup small CL swaps, by only fetching CL uptime accumulators if there is a tick crossing.
 * [#7562](https://github.com/osmosis-labs/osmosis/pull/7562) Speedup Protorev estimation logic by removing unnecessary taker fee simulations.
+* [#7595](https://github.com/osmosis-labs/osmosis/pull/7595) Fix cosmwasm pool model code ID migration.
 
 ### State Compatible
 

x/cosmwasmpool/gov.go

@@ -122,6 +122,10 @@ func (k Keeper) migrateCosmwasmPools(ctx sdk.Context, poolIds []uint64, newCodeI
 		if err != nil {
 			return err
 		}
+
+		// Update code ID to the updated one in state
+		cwPool.SetCodeId(newCodeId)
+		k.SetPool(ctx, cwPool)
 	}
 
 	// Whitelist new code id. No-op if already whitelisted.

x/cosmwasmpool/gov_test.go

@@ -147,6 +147,10 @@ func (s *CWPoolGovSuite) TestMigrateCosmwasmPools() {
 		preUploadCodeIdPlaceholder  uint64 = 1000
 		zeroCodeId                  uint64 = 0
 		defaultPoolCountToPreCreate uint64 = 3
+
+		// We create a code id for each pool. Since we provide code to upload in this test,
+		// we expect the code id to be one greater than the default pool count.
+		expectedNewCodeId uint64 = defaultPoolCountToPreCreate + 1
 	)
 
 	type MigrateMsg struct{}
@@ -181,7 +185,7 @@ func (s *CWPoolGovSuite) TestMigrateCosmwasmPools() {
 			byteCode:             emptyByteCode,
 			migrateMsg:           emptyMigrateMsg,
 
-			expectedCodeId: validCodeId,
+			expectedCodeId: expectedNewCodeId,
 		},
 		{
 			name:                                     "happy path with code id to upload",
@@ -192,7 +196,7 @@ func (s *CWPoolGovSuite) TestMigrateCosmwasmPools() {
 			migrateMsg:                               emptyMigrateMsg,
 			shouldWhitelistCWPoolModuleAccountUpload: true,
 
-			expectedCodeId: validCodeId,
+			expectedCodeId: expectedNewCodeId,
 		},
 		{
 			name:                                     "error: contract without migration entrypoint",
@@ -298,6 +302,15 @@ func (s *CWPoolGovSuite) TestMigrateCosmwasmPools() {
 			// Check that the code id is whitelisted.
 			s.Require().True(cosmwasmPoolKeeper.IsWhitelisted(s.Ctx, tc.expectedCodeId))
 
+			s.Require().NotEqual(0, len(tc.poolIdsToMigrate))
+			for _, poolID := range tc.poolIdsToMigrate {
+				// Check that the pool is migrated.
+				pool, err := cosmwasmPoolKeeper.GetPoolById(s.Ctx, poolID)
+				s.Require().NoError(err)
+
+				s.Require().Equal(tc.expectedCodeId, pool.GetCodeId())
+			}
+
 			// Validate that the event is emitted.
 			s.AssertEventEmitted(s.Ctx, types.TypeEvtMigratedCosmwasmPoolCode, 1)
 		})

x/cosmwasmpool/model/store_model.go

@@ -101,3 +101,8 @@ func (p CosmWasmPool) SetWasmKeeper(wasmKeeper types.WasmKeeper) {
 func (p *CosmWasmPool) GetPoolDenoms(ctx sdk.Context) []string {
 	panic("CosmWasmPool.GetPoolDenoms not implemented")
 }
+
+// SetCodeId implements types.CosmWasmExtension.
+func (p *CosmWasmPool) SetCodeId(codeId uint64) {
+	p.CodeId = codeId
+}

x/cosmwasmpool/types/pool.go

@@ -23,4 +23,6 @@ type CosmWasmExtension interface {
 	SetWasmKeeper(wasmKeeper WasmKeeper)
 
 	GetTotalPoolLiquidity(ctx sdk.Context) sdk.Coins
+
+	SetCodeId(codeId uint64)
 }