From 236d528ad24991878fcc03abd092264ef79888d1 Mon Sep 17 00:00:00 2001
From: PaddyMc <paddymchale@hotmail.com>
Date: Wed, 23 Oct 2024 10:30:34 +0100
Subject: [PATCH] fix: supply offsets for osmo token (#8751)

* fix: supply offsets for osmo token

* chore: add changelog entry
---
 CHANGELOG.md                      |  1 +
 app/upgrades/v27/constants.go     |  5 +-
 app/upgrades/v27/upgrades.go      | 10 ++++
 app/upgrades/v27/upgrades_test.go | 77 +++++++++++++++++++++++++++++++
 go.work                           |  1 +
 5 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 app/upgrades/v27/upgrades_test.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 51f8de50011..ee7796b4b4d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -46,6 +46,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * [#8682](https://github.com/osmosis-labs/osmosis/pull/8682) chore: bump cosmwasm-optimizer
 * [#8734](https://github.com/osmosis-labs/osmosis/pull/8734) chore: update cosmwasm vm
+* [#8751](https://github.com/osmosis-labs/osmosis/pull/8751) fix: supply offsets for osmo token
 * [#8764](https://github.com/osmosis-labs/osmosis/pull/8764) chore: add cosmwasm 1_3 feature
 * [#8779](https://github.com/osmosis-labs/osmosis/pull/8779) chore: bump cometbft/cosmos-sdk versions
 
diff --git a/app/upgrades/v27/constants.go b/app/upgrades/v27/constants.go
index 83764dc159d..598f223a0a1 100644
--- a/app/upgrades/v27/constants.go
+++ b/app/upgrades/v27/constants.go
@@ -7,7 +7,10 @@ import (
 )
 
 // UpgradeName defines the on-chain upgrade name for the Osmosis v27 upgrade.
-const UpgradeName = "v27"
+const (
+	UpgradeName = "v27"
+	OsmoToken   = "uosmo"
+)
 
 var Upgrade = upgrades.Upgrade{
 	UpgradeName:          UpgradeName,
diff --git a/app/upgrades/v27/upgrades.go b/app/upgrades/v27/upgrades.go
index 3b19e00e751..68c60ec3aec 100644
--- a/app/upgrades/v27/upgrades.go
+++ b/app/upgrades/v27/upgrades.go
@@ -23,6 +23,16 @@ func CreateUpgradeHandler(
 		if err != nil {
 			return nil, err
 		}
+		bk := keepers.BankKeeper
+
+		// Get the old offset from the old key: []byte{0x88}
+		offsetOld := bk.GetSupplyOffsetOld(ctx, OsmoToken)
+
+		// Add the old offset to the new key: collection.NewPrefix(88)
+		bk.AddSupplyOffset(ctx, OsmoToken, offsetOld)
+
+		// Remove the old key: []byte{0x88}
+		bk.RemoveOldSupplyOffset(ctx, OsmoToken)
 
 		return migrations, nil
 	}
diff --git a/app/upgrades/v27/upgrades_test.go b/app/upgrades/v27/upgrades_test.go
new file mode 100644
index 00000000000..33adcf70fdc
--- /dev/null
+++ b/app/upgrades/v27/upgrades_test.go
@@ -0,0 +1,77 @@
+package v27_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/suite"
+
+	"cosmossdk.io/core/appmodule"
+	"cosmossdk.io/core/header"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+
+	"github.com/osmosis-labs/osmosis/v26/app/apptesting"
+	v27 "github.com/osmosis-labs/osmosis/v26/app/upgrades/v27"
+
+	"cosmossdk.io/x/upgrade"
+
+	addresscodec "github.com/cosmos/cosmos-sdk/codec/address"
+
+	"github.com/osmosis-labs/osmosis/osmomath"
+)
+
+const (
+	v27UpgradeHeight = int64(10)
+)
+
+type UpgradeTestSuite struct {
+	apptesting.KeeperTestHelper
+	preModule appmodule.HasPreBlocker
+}
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+func (s *UpgradeTestSuite) TestUpgrade() {
+	s.Setup()
+	s.preModule = upgrade.NewAppModule(s.App.UpgradeKeeper, addresscodec.NewBech32Codec("osmo"))
+
+	s.PrepareSupplyOffsetTest()
+
+	// Run the upgrade
+	dummyUpgrade(s)
+	s.Require().NotPanics(func() {
+		_, err := s.preModule.PreBlock(s.Ctx)
+		s.Require().NoError(err)
+	})
+
+	s.ExecuteSupplyOffsetTest()
+}
+
+func dummyUpgrade(s *UpgradeTestSuite) {
+	s.Ctx = s.Ctx.WithBlockHeight(v27UpgradeHeight - 1)
+	plan := upgradetypes.Plan{Name: v27.Upgrade.UpgradeName, Height: v27UpgradeHeight}
+	err := s.App.UpgradeKeeper.ScheduleUpgrade(s.Ctx, plan)
+	s.Require().NoError(err)
+	_, err = s.App.UpgradeKeeper.GetUpgradePlan(s.Ctx)
+	s.Require().NoError(err)
+
+	s.Ctx = s.Ctx.WithHeaderInfo(header.Info{Height: v27UpgradeHeight, Time: s.Ctx.BlockTime().Add(time.Second)}).WithBlockHeight(v27UpgradeHeight)
+}
+
+func (s *UpgradeTestSuite) PrepareSupplyOffsetTest() {
+	// Set some supply offsets
+	s.App.BankKeeper.AddSupplyOffset(s.Ctx, v27.OsmoToken, osmomath.NewInt(1000))
+	s.App.BankKeeper.AddSupplyOffsetOld(s.Ctx, v27.OsmoToken, osmomath.NewInt(-500))
+}
+
+func (s *UpgradeTestSuite) ExecuteSupplyOffsetTest() {
+	coin := s.App.BankKeeper.GetSupplyWithOffset(s.Ctx, v27.OsmoToken)
+	offset := s.App.BankKeeper.GetSupplyOffset(s.Ctx, v27.OsmoToken)
+	oldOffset := s.App.BankKeeper.GetSupplyOffsetOld(s.Ctx, v27.OsmoToken)
+
+	s.Require().Equal("500uosmo", coin.String())
+	s.Require().Equal("500", offset.String())
+	s.Require().Equal("0", oldOffset.String())
+}
diff --git a/go.work b/go.work
index 7e6799a1ec3..6dd78cbdd8d 100644
--- a/go.work
+++ b/go.work
@@ -11,3 +11,4 @@ use ./x/ibc-hooks
 use ./x/epochs
 
 replace github.com/onsi/ginkgo => github.com/onsi/ginkgo v1.16.5
+replace	github.com/cosmos/cosmos-sdk => github.com/osmosis-labs/cosmos-sdk v0.50.10-v26-osmo-1.0.20241002161840-d2ff8f7253c7