added sttia rate limit and unit test

app/upgrades/v19/upgrades.go

@@ -3,6 +3,7 @@ package v19
 import (
  "time"
 
+ errorsmod "cosmossdk.io/errors"
  wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
  wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
  ratelimitkeeper "github.com/Stride-Labs/ibc-rate-limiting/ratelimit/keeper"
@@ -15,6 +16,11 @@ import (
 
 const (
  UpgradeName = "v19"
+
+ StTiaDenom = "stutia"
+ CelestiaTransferChannelId = "channel-162"
+ OsmosisTransferChannelId = "channel-5"
+ RateLimitDurationHours = 24
 )
 
 // CreateUpgradeHandler creates an SDK upgrade handler for v19
@@ -46,6 +52,11 @@ func CreateUpgradeHandler(
  return newVm, err
  }
 
+ // Add stTIA rate limits to Celestia and Osmosis
+ if err := AddStTiaRateLimit(ctx, ratelimitKeeper); err != nil {
+ return newVm, err
+ }
+
  return newVm, nil
  }
 }
@@ -65,3 +76,28 @@ func MigrateRateLimitModule(ctx sdk.Context, k ratelimitkeeper.Keeper) {
  hourEpoch.EpochStartHeight = ctx.BlockHeight()
  k.SetHourEpoch(ctx, hourEpoch)
 }
+
+// Add a 10% rate limit for stTIA to Celestia and Osmosis
+func AddStTiaRateLimit(ctx sdk.Context, k ratelimitkeeper.Keeper) error {
+ addRateLimitMsgTemplate := ratelimittypes.MsgAddRateLimit{
+ Denom: StTiaDenom,
+ MaxPercentSend: sdk.NewInt(10),
+ MaxPercentRecv: sdk.NewInt(10),
+ DurationHours: RateLimitDurationHours,
+ }
+
+ addCelestiaMsg := addRateLimitMsgTemplate
+ addCelestiaMsg.ChannelId = CelestiaTransferChannelId
+
+ addOsmosisMsg := addRateLimitMsgTemplate
+ addOsmosisMsg.ChannelId = OsmosisTransferChannelId
+
+ if err := k.AddRateLimit(ctx, &addCelestiaMsg); err != nil {
+ return errorsmod.Wrapf(err, "unable to add stTIA rate limit to celestia")
+ }
+ if err := k.AddRateLimit(ctx, &addOsmosisMsg); err != nil {
+ return errorsmod.Wrapf(err, "unable to add stTIA rate limit to celestia")
+ }
+
+ return nil
+}

app/upgrades/v19/upgrades_test.go

@@ -4,16 +4,24 @@ import (
  "testing"
 
  sdkmath "cosmossdk.io/math"
+ sdk "github.com/cosmos/cosmos-sdk/types"
+ transfertypes "github.com/cosmos/ibc-go/v7/modules/apps/transfer/types"
+ channeltypes "github.com/cosmos/ibc-go/v7/modules/core/04-channel/types"
 
  ratelimittypes "github.com/Stride-Labs/ibc-rate-limiting/ratelimit/types"
  "github.com/cosmos/cosmos-sdk/store/prefix"
  "github.com/stretchr/testify/suite"
 
  "github.com/Stride-Labs/stride/v18/app"
  "github.com/Stride-Labs/stride/v18/app/apptesting"
+ v19 "github.com/Stride-Labs/stride/v18/app/upgrades/v19"
  legacyratelimittypes "github.com/Stride-Labs/stride/v18/app/upgrades/v19/legacyratelimit/types"
 )
 
+var (
+ StTiaSupply = sdkmath.NewInt(1000)
+)
+
 type UpgradeTestSuite struct {
  apptesting.AppTestHelper
 }
@@ -29,15 +37,19 @@ func (s *UpgradeTestSuite) SetupTest() {
 func (s *UpgradeTestSuite) TestUpgrade() {
  dummyUpgradeHeight := int64(5)
 
- checkRateLimitsAfterUpgrade := s.SetupRateLimitsBeforeUpgrade()
+ // Setup state before upgrade
+ checkMigratedRateLimits := s.SetupRateLimitMigration()
+ checkStTiaRateLimits := s.SetupStTiaRateLimits()
 
  // Run through upgrade
  s.ConfirmUpgradeSucceededs("v19", dummyUpgradeHeight)
 
- checkRateLimitsAfterUpgrade()
+ // Check state after upgrade
+ checkMigratedRateLimits()
+ checkStTiaRateLimits()
 }
 
-func (s *UpgradeTestSuite) SetupRateLimitsBeforeUpgrade() func() {
+func (s *UpgradeTestSuite) SetupRateLimitMigration() func() {
  rateLimitStore := s.Ctx.KVStore(s.App.GetKey(ratelimittypes.StoreKey))
  cdc := app.MakeEncodingConfig().Marshaler
 
@@ -86,9 +98,64 @@ func (s *UpgradeTestSuite) SetupRateLimitsBeforeUpgrade() func() {
  hostzoneStore := prefix.NewStore(rateLimitStore, ratelimittypes.RateLimitKeyPrefix)
  hostzoneStore.Set(ratelimittypes.GetRateLimitItemKey(denom, channelId), initialRateLimitBz)
 
+ // Return a callback to check the state after the upgrade
  return func() {
  actualRateLimit, found := s.App.RatelimitKeeper.GetRateLimit(s.Ctx, denom, channelId)
  s.Require().True(found, "rate limit should have been found")
  s.Require().Equal(expectedRateLimit, actualRateLimit, "rate limit after upgrade")
  }
 }
+
+func (s *UpgradeTestSuite) SetupStTiaRateLimits() func() {
+ // mint sttia so that there is a channel value
+ s.FundAccount(s.TestAccs[0], sdk.NewCoin(v19.StTiaDenom, StTiaSupply))
+
+ // Mock out a channel for osmosis and celstia
+ s.App.IBCKeeper.ChannelKeeper.SetChannel(s.Ctx, transfertypes.PortID, v19.CelestiaTransferChannelId, channeltypes.Channel{})
+ s.App.IBCKeeper.ChannelKeeper.SetChannel(s.Ctx, transfertypes.PortID, v19.OsmosisTransferChannelId, channeltypes.Channel{})
+
+ // Return a callback to check the rate limits were added after the upgrade
+ return func() {
+ expectedRateLimitToCelestia := ratelimittypes.RateLimit{
+ Path: &ratelimittypes.Path{
+ Denom: v19.StTiaDenom,
+ ChannelId: v19.CelestiaTransferChannelId,
+ },
+ Flow: &ratelimittypes.Flow{
+ Inflow: sdkmath.NewInt(0),
+ Outflow: sdkmath.NewInt(0),
+ ChannelValue: StTiaSupply,
+ },
+ Quota: &ratelimittypes.Quota{
+ MaxPercentSend: sdkmath.NewInt(10),
+ MaxPercentRecv: sdkmath.NewInt(10),
+ DurationHours: 24,
+ },
+ }
+
+ expectedRateLimitToOsmosis := ratelimittypes.RateLimit{
+ Path: &ratelimittypes.Path{
+ Denom: v19.StTiaDenom,
+ ChannelId: v19.OsmosisTransferChannelId,
+ },
+ Flow: &ratelimittypes.Flow{
+ Inflow: sdkmath.NewInt(0),
+ Outflow: sdkmath.NewInt(0),
+ ChannelValue: StTiaSupply,
+ },
+ Quota: &ratelimittypes.Quota{
+ MaxPercentSend: sdkmath.NewInt(10),
+ MaxPercentRecv: sdkmath.NewInt(10),
+ DurationHours: 24,
+ },
+ }
+
+ actualCelestiaRateLimit, found := s.App.RatelimitKeeper.GetRateLimit(s.Ctx, v19.StTiaDenom, v19.CelestiaTransferChannelId)
+ s.Require().True(found, "rate limit to celestia should have been found")
+ s.Require().Equal(expectedRateLimitToCelestia, actualCelestiaRateLimit)
+
+ actualOsmosisRateLimit, found := s.App.RatelimitKeeper.GetRateLimit(s.Ctx, v19.StTiaDenom, v19.OsmosisTransferChannelId)
+ s.Require().True(found, "rate limit to osmosis should have been found")
+ s.Require().Equal(expectedRateLimitToOsmosis, actualOsmosisRateLimit)
+ }
+}