refactor: remove PokePool from the PoolI interface, define on extension instead (#3035)

* refactor: remove PokePool from the PoolI interface, use extension instead

* changelog

* fix weights bug

* Update x/gamm/keeper/pool_test.go

Co-authored-by: Matt, Park <45252226+mattverse@users.noreply.github.com>

* Update x/gamm/keeper/pool_test.go

* address comments

Co-authored-by: Matt, Park <45252226+mattverse@users.noreply.github.com>

CHANGELOG.md

@@ -55,6 +55,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Misc Improvements
 
 * [#2804](https://github.com/osmosis-labs/osmosis/pull/2804) Improve error handling and messages when parsing pool assets.
+* [#3035](https://github.com/osmosis-labs/osmosis/pull/3035) Remove `PokePool` from `PoolI` interface. Define on a new WeightedPoolExtension` instead.
 
 ## v12.0.0
 

tests/mocks/mock_pool.go

 // Code generated by MockGen. DO NOT EDIT.
 // Source: x/gamm/types/pool.go
 
-// Package mocks is a generated GoMock package.
+// Package mock_types is a generated GoMock package.
 package mocks
 
 import (
@@ -255,18 +255,6 @@ func (mr *MockPoolIMockRecorder) JoinPoolNoSwap(ctx, tokensIn, swapFee interface
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "JoinPoolNoSwap", reflect.TypeOf((*MockPoolI)(nil).JoinPoolNoSwap), ctx, tokensIn, swapFee)
 }
 
-// PokePool mocks base method.
-func (m *MockPoolI) PokePool(blockTime time.Time) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "PokePool", blockTime)
-}
-
-// PokePool indicates an expected call of PokePool.
-func (mr *MockPoolIMockRecorder) PokePool(blockTime interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PokePool", reflect.TypeOf((*MockPoolI)(nil).PokePool), blockTime)
-}
-
 // ProtoMessage mocks base method.
 func (m *MockPoolI) ProtoMessage() {
 	m.ctrl.T.Helper()
@@ -650,18 +638,6 @@ func (mr *MockPoolAmountOutExtensionMockRecorder) JoinPoolTokenInMaxShareAmountO
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "JoinPoolTokenInMaxShareAmountOut", reflect.TypeOf((*MockPoolAmountOutExtension)(nil).JoinPoolTokenInMaxShareAmountOut), ctx, tokenInDenom, shareOutAmount)
 }
 
-// PokePool mocks base method.
-func (m *MockPoolAmountOutExtension) PokePool(blockTime time.Time) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "PokePool", blockTime)
-}
-
-// PokePool indicates an expected call of PokePool.
-func (mr *MockPoolAmountOutExtensionMockRecorder) PokePool(blockTime interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PokePool", reflect.TypeOf((*MockPoolAmountOutExtension)(nil).PokePool), blockTime)
-}
-
 // ProtoMessage mocks base method.
 func (m *MockPoolAmountOutExtension) ProtoMessage() {
 	m.ctrl.T.Helper()
@@ -744,3 +720,341 @@ func (mr *MockPoolAmountOutExtensionMockRecorder) SwapOutAmtGivenIn(ctx, tokenIn
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SwapOutAmtGivenIn", reflect.TypeOf((*MockPoolAmountOutExtension)(nil).SwapOutAmtGivenIn), ctx, tokenIn, tokenOutDenom, swapFee)
 }
+
+// MockPokePoolExtension is a mock of PokePoolExtension interface.
+type MockPokePoolExtension struct {
+	ctrl     *gomock.Controller
+	recorder *MockPokePoolExtensionMockRecorder
+}
+
+// MockPokePoolExtensionMockRecorder is the mock recorder for MockPokePoolExtension.
+type MockPokePoolExtensionMockRecorder struct {
+	mock *MockPokePoolExtension
+}
+
+// NewMockPokePoolExtension creates a new mock instance.
+func NewMockPokePoolExtension(ctrl *gomock.Controller) *MockPokePoolExtension {
+	mock := &MockPokePoolExtension{ctrl: ctrl}
+	mock.recorder = &MockPokePoolExtensionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockPokePoolExtension) EXPECT() *MockPokePoolExtensionMockRecorder {
+	return m.recorder
+}
+
+// CalcExitPoolCoinsFromShares mocks base method.
+func (m *MockPokePoolExtension) CalcExitPoolCoinsFromShares(ctx types.Context, numShares types.Int, exitFee types.Dec) (types.Coins, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CalcExitPoolCoinsFromShares", ctx, numShares, exitFee)
+	ret0, _ := ret[0].(types.Coins)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CalcExitPoolCoinsFromShares indicates an expected call of CalcExitPoolCoinsFromShares.
+func (mr *MockPokePoolExtensionMockRecorder) CalcExitPoolCoinsFromShares(ctx, numShares, exitFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalcExitPoolCoinsFromShares", reflect.TypeOf((*MockPokePoolExtension)(nil).CalcExitPoolCoinsFromShares), ctx, numShares, exitFee)
+}
+
+// CalcInAmtGivenOut mocks base method.
+func (m *MockPokePoolExtension) CalcInAmtGivenOut(ctx types.Context, tokenOut types.Coins, tokenInDenom string, swapFee types.Dec) (types.Coin, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CalcInAmtGivenOut", ctx, tokenOut, tokenInDenom, swapFee)
+	ret0, _ := ret[0].(types.Coin)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CalcInAmtGivenOut indicates an expected call of CalcInAmtGivenOut.
+func (mr *MockPokePoolExtensionMockRecorder) CalcInAmtGivenOut(ctx, tokenOut, tokenInDenom, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalcInAmtGivenOut", reflect.TypeOf((*MockPokePoolExtension)(nil).CalcInAmtGivenOut), ctx, tokenOut, tokenInDenom, swapFee)
+}
+
+// CalcJoinPoolNoSwapShares mocks base method.
+func (m *MockPokePoolExtension) CalcJoinPoolNoSwapShares(ctx types.Context, tokensIn types.Coins, swapFee types.Dec) (types.Int, types.Coins, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CalcJoinPoolNoSwapShares", ctx, tokensIn, swapFee)
+	ret0, _ := ret[0].(types.Int)
+	ret1, _ := ret[1].(types.Coins)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
+}
+
+// CalcJoinPoolNoSwapShares indicates an expected call of CalcJoinPoolNoSwapShares.
+func (mr *MockPokePoolExtensionMockRecorder) CalcJoinPoolNoSwapShares(ctx, tokensIn, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalcJoinPoolNoSwapShares", reflect.TypeOf((*MockPokePoolExtension)(nil).CalcJoinPoolNoSwapShares), ctx, tokensIn, swapFee)
+}
+
+// CalcJoinPoolShares mocks base method.
+func (m *MockPokePoolExtension) CalcJoinPoolShares(ctx types.Context, tokensIn types.Coins, swapFee types.Dec) (types.Int, types.Coins, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CalcJoinPoolShares", ctx, tokensIn, swapFee)
+	ret0, _ := ret[0].(types.Int)
+	ret1, _ := ret[1].(types.Coins)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
+}
+
+// CalcJoinPoolShares indicates an expected call of CalcJoinPoolShares.
+func (mr *MockPokePoolExtensionMockRecorder) CalcJoinPoolShares(ctx, tokensIn, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalcJoinPoolShares", reflect.TypeOf((*MockPokePoolExtension)(nil).CalcJoinPoolShares), ctx, tokensIn, swapFee)
+}
+
+// CalcOutAmtGivenIn mocks base method.
+func (m *MockPokePoolExtension) CalcOutAmtGivenIn(ctx types.Context, tokenIn types.Coins, tokenOutDenom string, swapFee types.Dec) (types.Coin, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CalcOutAmtGivenIn", ctx, tokenIn, tokenOutDenom, swapFee)
+	ret0, _ := ret[0].(types.Coin)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CalcOutAmtGivenIn indicates an expected call of CalcOutAmtGivenIn.
+func (mr *MockPokePoolExtensionMockRecorder) CalcOutAmtGivenIn(ctx, tokenIn, tokenOutDenom, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalcOutAmtGivenIn", reflect.TypeOf((*MockPokePoolExtension)(nil).CalcOutAmtGivenIn), ctx, tokenIn, tokenOutDenom, swapFee)
+}
+
+// ExitPool mocks base method.
+func (m *MockPokePoolExtension) ExitPool(ctx types.Context, numShares types.Int, exitFee types.Dec) (types.Coins, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ExitPool", ctx, numShares, exitFee)
+	ret0, _ := ret[0].(types.Coins)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ExitPool indicates an expected call of ExitPool.
+func (mr *MockPokePoolExtensionMockRecorder) ExitPool(ctx, numShares, exitFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExitPool", reflect.TypeOf((*MockPokePoolExtension)(nil).ExitPool), ctx, numShares, exitFee)
+}
+
+// GetAddress mocks base method.
+func (m *MockPokePoolExtension) GetAddress() types.AccAddress {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetAddress")
+	ret0, _ := ret[0].(types.AccAddress)
+	return ret0
+}
+
+// GetAddress indicates an expected call of GetAddress.
+func (mr *MockPokePoolExtensionMockRecorder) GetAddress() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAddress", reflect.TypeOf((*MockPokePoolExtension)(nil).GetAddress))
+}
+
+// GetExitFee mocks base method.
+func (m *MockPokePoolExtension) GetExitFee(ctx types.Context) types.Dec {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetExitFee", ctx)
+	ret0, _ := ret[0].(types.Dec)
+	return ret0
+}
+
+// GetExitFee indicates an expected call of GetExitFee.
+func (mr *MockPokePoolExtensionMockRecorder) GetExitFee(ctx interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetExitFee", reflect.TypeOf((*MockPokePoolExtension)(nil).GetExitFee), ctx)
+}
+
+// GetId mocks base method.
+func (m *MockPokePoolExtension) GetId() uint64 {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetId")
+	ret0, _ := ret[0].(uint64)
+	return ret0
+}
+
+// GetId indicates an expected call of GetId.
+func (mr *MockPokePoolExtensionMockRecorder) GetId() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetId", reflect.TypeOf((*MockPokePoolExtension)(nil).GetId))
+}
+
+// GetSwapFee mocks base method.
+func (m *MockPokePoolExtension) GetSwapFee(ctx types.Context) types.Dec {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetSwapFee", ctx)
+	ret0, _ := ret[0].(types.Dec)
+	return ret0
+}
+
+// GetSwapFee indicates an expected call of GetSwapFee.
+func (mr *MockPokePoolExtensionMockRecorder) GetSwapFee(ctx interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSwapFee", reflect.TypeOf((*MockPokePoolExtension)(nil).GetSwapFee), ctx)
+}
+
+// GetTotalPoolLiquidity mocks base method.
+func (m *MockPokePoolExtension) GetTotalPoolLiquidity(ctx types.Context) types.Coins {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetTotalPoolLiquidity", ctx)
+	ret0, _ := ret[0].(types.Coins)
+	return ret0
+}
+
+// GetTotalPoolLiquidity indicates an expected call of GetTotalPoolLiquidity.
+func (mr *MockPokePoolExtensionMockRecorder) GetTotalPoolLiquidity(ctx interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTotalPoolLiquidity", reflect.TypeOf((*MockPokePoolExtension)(nil).GetTotalPoolLiquidity), ctx)
+}
+
+// GetTotalShares mocks base method.
+func (m *MockPokePoolExtension) GetTotalShares() types.Int {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetTotalShares")
+	ret0, _ := ret[0].(types.Int)
+	return ret0
+}
+
+// GetTotalShares indicates an expected call of GetTotalShares.
+func (mr *MockPokePoolExtensionMockRecorder) GetTotalShares() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTotalShares", reflect.TypeOf((*MockPokePoolExtension)(nil).GetTotalShares))
+}
+
+// IsActive mocks base method.
+func (m *MockPokePoolExtension) IsActive(ctx types.Context) bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "IsActive", ctx)
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// IsActive indicates an expected call of IsActive.
+func (mr *MockPokePoolExtensionMockRecorder) IsActive(ctx interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsActive", reflect.TypeOf((*MockPokePoolExtension)(nil).IsActive), ctx)
+}
+
+// JoinPool mocks base method.
+func (m *MockPokePoolExtension) JoinPool(ctx types.Context, tokensIn types.Coins, swapFee types.Dec) (types.Int, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "JoinPool", ctx, tokensIn, swapFee)
+	ret0, _ := ret[0].(types.Int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// JoinPool indicates an expected call of JoinPool.
+func (mr *MockPokePoolExtensionMockRecorder) JoinPool(ctx, tokensIn, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "JoinPool", reflect.TypeOf((*MockPokePoolExtension)(nil).JoinPool), ctx, tokensIn, swapFee)
+}
+
+// JoinPoolNoSwap mocks base method.
+func (m *MockPokePoolExtension) JoinPoolNoSwap(ctx types.Context, tokensIn types.Coins, swapFee types.Dec) (types.Int, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "JoinPoolNoSwap", ctx, tokensIn, swapFee)
+	ret0, _ := ret[0].(types.Int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// JoinPoolNoSwap indicates an expected call of JoinPoolNoSwap.
+func (mr *MockPokePoolExtensionMockRecorder) JoinPoolNoSwap(ctx, tokensIn, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "JoinPoolNoSwap", reflect.TypeOf((*MockPokePoolExtension)(nil).JoinPoolNoSwap), ctx, tokensIn, swapFee)
+}
+
+// PokePool mocks base method.
+func (m *MockPokePoolExtension) PokePool(blockTime time.Time) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "PokePool", blockTime)
+}
+
+// PokePool indicates an expected call of PokePool.
+func (mr *MockPokePoolExtensionMockRecorder) PokePool(blockTime interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PokePool", reflect.TypeOf((*MockPokePoolExtension)(nil).PokePool), blockTime)
+}
+
+// ProtoMessage mocks base method.
+func (m *MockPokePoolExtension) ProtoMessage() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "ProtoMessage")
+}
+
+// ProtoMessage indicates an expected call of ProtoMessage.
+func (mr *MockPokePoolExtensionMockRecorder) ProtoMessage() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProtoMessage", reflect.TypeOf((*MockPokePoolExtension)(nil).ProtoMessage))
+}
+
+// Reset mocks base method.
+func (m *MockPokePoolExtension) Reset() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Reset")
+}
+
+// Reset indicates an expected call of Reset.
+func (mr *MockPokePoolExtensionMockRecorder) Reset() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Reset", reflect.TypeOf((*MockPokePoolExtension)(nil).Reset))
+}
+
+// SpotPrice mocks base method.
+func (m *MockPokePoolExtension) SpotPrice(ctx types.Context, baseAssetDenom, quoteAssetDenom string) (types.Dec, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SpotPrice", ctx, baseAssetDenom, quoteAssetDenom)
+	ret0, _ := ret[0].(types.Dec)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SpotPrice indicates an expected call of SpotPrice.
+func (mr *MockPokePoolExtensionMockRecorder) SpotPrice(ctx, baseAssetDenom, quoteAssetDenom interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SpotPrice", reflect.TypeOf((*MockPokePoolExtension)(nil).SpotPrice), ctx, baseAssetDenom, quoteAssetDenom)
+}
+
+// String mocks base method.
+func (m *MockPokePoolExtension) String() string {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "String")
+	ret0, _ := ret[0].(string)
+	return ret0
+}
+
+// String indicates an expected call of String.
+func (mr *MockPokePoolExtensionMockRecorder) String() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "String", reflect.TypeOf((*MockPokePoolExtension)(nil).String))
+}
+
+// SwapInAmtGivenOut mocks base method.
+func (m *MockPokePoolExtension) SwapInAmtGivenOut(ctx types.Context, tokenOut types.Coins, tokenInDenom string, swapFee types.Dec) (types.Coin, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SwapInAmtGivenOut", ctx, tokenOut, tokenInDenom, swapFee)
+	ret0, _ := ret[0].(types.Coin)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SwapInAmtGivenOut indicates an expected call of SwapInAmtGivenOut.
+func (mr *MockPokePoolExtensionMockRecorder) SwapInAmtGivenOut(ctx, tokenOut, tokenInDenom, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SwapInAmtGivenOut", reflect.TypeOf((*MockPokePoolExtension)(nil).SwapInAmtGivenOut), ctx, tokenOut, tokenInDenom, swapFee)
+}
+
+// SwapOutAmtGivenIn mocks base method.
+func (m *MockPokePoolExtension) SwapOutAmtGivenIn(ctx types.Context, tokenIn types.Coins, tokenOutDenom string, swapFee types.Dec) (types.Coin, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SwapOutAmtGivenIn", ctx, tokenIn, tokenOutDenom, swapFee)
+	ret0, _ := ret[0].(types.Coin)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SwapOutAmtGivenIn indicates an expected call of SwapOutAmtGivenIn.
+func (mr *MockPokePoolExtensionMockRecorder) SwapOutAmtGivenIn(ctx, tokenIn, tokenOutDenom, swapFee interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SwapOutAmtGivenIn", reflect.TypeOf((*MockPokePoolExtension)(nil).SwapOutAmtGivenIn), ctx, tokenIn, tokenOutDenom, swapFee)
+}

x/gamm/keeper/invariants.go

@@ -12,7 +12,7 @@ import (
 
 const poolBalanceInvariantName = "pool-account-balance-equals-expected"
 
-// RegisterInvariants registers all governance invariants.
+// RegisterInvariants registers all gamm invariants.
 func RegisterInvariants(ir sdk.InvariantRegistry, keeper Keeper, bk types.BankKeeper) {
 	ir.RegisterRoute(types.ModuleName, poolBalanceInvariantName, PoolAccountInvariant(keeper, bk))
 }

x/gamm/keeper/keeper_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/osmosis-labs/osmosis/v12/app/apptesting"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
 	balancertypes "github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/stableswap"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
 )
 
@@ -34,9 +35,7 @@ func (suite *KeeperTestSuite) prepareCustomBalancerPool(
 	poolAssets []balancertypes.PoolAsset,
 	poolParams balancer.PoolParams,
 ) uint64 {
-	for _, acc := range suite.TestAccs {
-		suite.FundAcc(acc, balances)
-	}
+	suite.fundAllAccountsWith(balances)
 
 	poolID, err := suite.App.GAMMKeeper.CreatePool(
 		suite.Ctx,
@@ -46,3 +45,26 @@ func (suite *KeeperTestSuite) prepareCustomBalancerPool(
 
 	return poolID
 }
+
+func (suite *KeeperTestSuite) prepareCustomStableswapPool(
+	balances sdk.Coins,
+	poolParams stableswap.PoolParams,
+	initialLiquidity sdk.Coins,
+	scalingFactors []uint64,
+) uint64 {
+	suite.fundAllAccountsWith(balances)
+
+	poolID, err := suite.App.GAMMKeeper.CreatePool(
+		suite.Ctx,
+		stableswap.NewMsgCreateStableswapPool(suite.TestAccs[0], poolParams, initialLiquidity, scalingFactors, ""),
+	)
+	suite.Require().NoError(err)
+
+	return poolID
+}
+
+func (suite *KeeperTestSuite) fundAllAccountsWith(balances sdk.Coins) {
+	for _, acc := range suite.TestAccs {
+		suite.FundAcc(acc, balances)
+	}
+}

x/gamm/keeper/pool.go

@@ -22,8 +22,8 @@ func (k Keeper) UnmarshalPool(bz []byte) (types.PoolI, error) {
 	return acc, k.cdc.UnmarshalInterface(bz, &acc)
 }
 
-// GetPoolAndPoke returns a PoolI based on it's identifier if one exists. Prior
-// to returning the pool, the weights of the pool are updated via PokePool.
+// GetPoolAndPoke returns a PoolI based on it's identifier if one exists. If poolId corresponds
+// to a pool with weights (e.g. balancer), the weights of the pool are updated via PokePool prior to returning.
 // TODO: Consider rename to GetPool due to downstream API confusion.
 func (k Keeper) GetPoolAndPoke(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
 	store := ctx.KVStore(k.storeKey)
@@ -39,7 +39,9 @@ func (k Keeper) GetPoolAndPoke(ctx sdk.Context, poolId uint64) (types.PoolI, err
 		return nil, err
 	}
 
-	pool.PokePool(ctx.BlockTime())
+	if pokePool, ok := pool.(types.WeightedPoolExtension); ok {
+		pokePool.PokePool(ctx.BlockTime())
+	}
 
 	return pool, nil
 }
@@ -74,7 +76,9 @@ func (k Keeper) GetPoolsAndPoke(ctx sdk.Context) (res []types.PoolI, err error)
 			return nil, err
 		}
 
-		pool.PokePool(ctx.BlockTime())
+		if pokePool, ok := pool.(types.WeightedPoolExtension); ok {
+			pokePool.PokePool(ctx.BlockTime())
+		}
 		res = append(res, pool)
 	}
 

x/gamm/keeper/pool_test.go

 package keeper_test
 
+import (
+	"time"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	balancertypes "github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/stableswap"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
+)
+
 // import (
 // 	"math/rand"
 // 	"time"
@@ -220,3 +231,87 @@ package keeper_test
 // 			"Expected equal %s: %d, %d", amt.Denom, amt.Amount.Int64(), sdk.NewInt(1000).Int64())
 // 	}
 // }
+
+// TestGetPoolAndPoke tests that the right pools is returned from GetPoolAndPoke.
+// For the pools implementing the weighted extension, asserts that PokePool is called.
+func (suite *KeeperTestSuite) TestGetPoolAndPoke() {
+	const (
+		startTime = 1000
+		blockTime = startTime + 100
+	)
+
+	// N.B.: We make a copy because SmoothWeightChangeParams get mutated.
+	// We would like to avoid mutating global pool assets that are used in other tests.
+	defaultPoolAssetsCopy := make([]balancertypes.PoolAsset, 2)
+	copy(defaultPoolAssetsCopy, defaultPoolAssets)
+
+	startPoolWeightAssets := []balancertypes.PoolAsset{
+		{
+			Weight: defaultPoolAssets[0].Weight.Quo(sdk.NewInt(2)),
+			Token:  defaultPoolAssets[0].Token,
+		},
+		{
+			Weight: defaultPoolAssets[1].Weight.Mul(sdk.NewInt(3)),
+			Token:  defaultPoolAssets[1].Token,
+		},
+	}
+
+	tests := map[string]struct {
+		isPokePool bool
+		poolId     uint64
+	}{
+		"weighted pool - change weights": {
+			isPokePool: true,
+			poolId: suite.prepareCustomBalancerPool(defaultAcctFunds, startPoolWeightAssets, balancer.PoolParams{
+				SwapFee: defaultSwapFee,
+				ExitFee: defaultExitFee,
+				SmoothWeightChangeParams: &balancer.SmoothWeightChangeParams{
+					StartTime:          time.Unix(startTime, 0), // start time is before block time so the weights should change
+					Duration:           time.Hour,
+					InitialPoolWeights: startPoolWeightAssets,
+					TargetPoolWeights:  defaultPoolAssetsCopy,
+				},
+			}),
+		},
+		"non weighted pool": {
+			poolId: suite.prepareCustomStableswapPool(
+				defaultAcctFunds,
+				stableswap.PoolParams{
+					SwapFee: defaultSwapFee,
+					ExitFee: defaultExitFee,
+				},
+				sdk.NewCoins(sdk.NewCoin(defaultAcctFunds[0].Denom, defaultAcctFunds[0].Amount.QuoRaw(2)), sdk.NewCoin(defaultAcctFunds[1].Denom, defaultAcctFunds[1].Amount.QuoRaw(2))),
+				[]uint64{1, 1},
+			)},
+	}
+
+	for name, tc := range tests {
+		suite.Run(name, func() {
+			k := suite.App.GAMMKeeper
+			ctx := suite.Ctx.WithBlockTime(time.Unix(blockTime, 0))
+
+			pool, err := k.GetPoolAndPoke(ctx, tc.poolId)
+
+			suite.Require().NoError(err)
+			suite.Require().Equal(tc.poolId, pool.GetId())
+
+			if tc.isPokePool {
+				pokePool, ok := pool.(types.WeightedPoolExtension)
+				suite.Require().True(ok)
+
+				poolAssetWeight0, err := pokePool.GetTokenWeight(startPoolWeightAssets[0].Token.Denom)
+				suite.Require().NoError(err)
+
+				poolAssetWeight1, err := pokePool.GetTokenWeight(startPoolWeightAssets[1].Token.Denom)
+				suite.Require().NoError(err)
+
+				suite.Require().NotEqual(startPoolWeightAssets[0].Weight, poolAssetWeight0)
+				suite.Require().NotEqual(startPoolWeightAssets[1].Weight, poolAssetWeight1)
+				return
+			}
+
+			_, ok := pool.(types.WeightedPoolExtension)
+			suite.Require().False(ok)
+		})
+	}
+}

x/gamm/module.go

@@ -29,6 +29,7 @@ import (
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/client/cli"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/keeper"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/stableswap"
 	simulation "github.com/osmosis-labs/osmosis/v12/x/gamm/simulation"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
 )
@@ -88,7 +89,7 @@ func (b AppModuleBasic) GetQueryCmd() *cobra.Command {
 func (AppModuleBasic) RegisterInterfaces(registry codectypes.InterfaceRegistry) {
 	types.RegisterInterfaces(registry)
 	balancer.RegisterInterfaces(registry)
-	// stableswap.RegisterInterfaces(registry)
+	stableswap.RegisterInterfaces(registry)
 }
 
 type AppModule struct {

x/gamm/pool-models/balancer/pool.go

@@ -29,6 +29,7 @@ const (
 var (
 	_ types.PoolI                  = &Pool{}
 	_ types.PoolAmountOutExtension = &Pool{}
+	_ types.WeightedPoolExtension  = &Pool{}
 )
 
 // NewPool returns a weighted CPMM pool with the provided parameters, and initial assets.

x/gamm/pool-models/stableswap/pool.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -317,9 +316,6 @@ func (p Pool) CalcExitPoolCoinsFromShares(ctx sdk.Context, exitingShares sdk.Int
 	return cfmm_common.CalcExitPool(ctx, &p, exitingShares, exitFee)
 }
 
-// no-op for stableswap
-func (p *Pool) PokePool(blockTime time.Time) {}
-
 // SetStableSwapScalingFactors sets scaling factors for pool to the given amount
 // It should only be able to be successfully called by the pool's ScalingFactorGovernor
 // TODO: move commented test for this function from x/gamm/keeper/pool_service_test.go once a pool_test.go file has been created for stableswap

x/gamm/types/pool.go

@@ -79,10 +79,6 @@ type PoolI interface {
 	// CalcExitPoolCoinsFromShares returns how many coins ExitPool would return on these arguments.
 	// This does not mutate the pool, or state.
 	CalcExitPoolCoinsFromShares(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error)
-
-	// PokePool determines if a pool's weights need to be updated and updates
-	// them if so.
-	PokePool(blockTime time.Time)
 }
 
 // PoolAmountOutExtension is an extension of the PoolI
@@ -126,6 +122,19 @@ type PoolAmountOutExtension interface {
 	IncreaseLiquidity(sharesOut sdk.Int, coinsIn sdk.Coins)
 }
 
+// WeightedPoolExtension is an extension of the PoolI interface
+// That defines an additional API for handling the pool's weights.
+type WeightedPoolExtension interface {
+	PoolI
+
+	// PokePool determines if a pool's weights need to be updated and updates
+	// them if so.
+	PokePool(blockTime time.Time)
+
+	// GetTokenWeight returns the weight of the specified token in the pool.
+	GetTokenWeight(denom string) (sdk.Int, error)
+}
+
 func NewPoolAddress(poolId uint64) sdk.AccAddress {
 	key := append([]byte("pool"), sdk.Uint64ToBigEndian(poolId)...)
 	return address.Module(ModuleName, key)