From 6072f0bb1d92172c7e3b5ac1e8d9178ae3b7f262 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C4=90=C3=B4ng=20Li=E1=BB=81u?=
 <93205232+DongLieu@users.noreply.github.com>
Date: Wed, 11 Oct 2023 19:25:22 +0700
Subject: [PATCH] Test for GetChainIDFromContext (#685)

* Test for GetChainIDFromContext

* lint
---
 x/interchainstaking/keeper/keeper_test.go | 52 +++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/x/interchainstaking/keeper/keeper_test.go b/x/interchainstaking/keeper/keeper_test.go
index 5361e2a71..1ad9a87bb 100644
--- a/x/interchainstaking/keeper/keeper_test.go
+++ b/x/interchainstaking/keeper/keeper_test.go
@@ -1,6 +1,8 @@
 package keeper_test
 
 import (
+	"context"
+	"errors"
 	"testing"
 	"time"
 
@@ -20,9 +22,11 @@ import (
 	ibctesting "github.com/cosmos/ibc-go/v5/testing"
 
 	"github.com/quicksilver-zone/quicksilver/app"
+	"github.com/quicksilver-zone/quicksilver/utils"
 	"github.com/quicksilver-zone/quicksilver/utils/addressutils"
 	"github.com/quicksilver-zone/quicksilver/utils/randomutils"
 	ics "github.com/quicksilver-zone/quicksilver/x/interchainstaking"
+	interchainstakingkeeper "github.com/quicksilver-zone/quicksilver/x/interchainstaking/keeper"
 	icstypes "github.com/quicksilver-zone/quicksilver/x/interchainstaking/types"
 )
 
@@ -665,6 +669,54 @@ func (suite *KeeperTestSuite) TestOverrideRedemptionRateNoCap() {
 	suite.Equal(sdk.NewDecWithPrec(676666666666666667, 18), zone.RedemptionRate)
 }
 
+func (suite *KeeperTestSuite) TestGetChainIDFromContext() {
+	testCase := []struct {
+		name            string
+		setup           func() (*interchainstakingkeeper.Keeper, sdk.Context)
+		wantErr         bool
+		expectedErr     error
+		expectedChainID string
+	}{
+		{
+			name: "connectionID not in context",
+			setup: func() (*interchainstakingkeeper.Keeper, sdk.Context) {
+				suite.SetupTest()
+				suite.setupTestZones()
+				return suite.GetQuicksilverApp(suite.chainA).InterchainstakingKeeper, suite.chainA.GetContext()
+			},
+			wantErr:     true,
+			expectedErr: errors.New("connectionID not in context"),
+		},
+		{
+			name: "get chainID success",
+			setup: func() (*interchainstakingkeeper.Keeper, sdk.Context) {
+				suite.SetupTest()
+				suite.setupTestZones()
+				ctx := suite.chainA.GetContext()
+
+				ctx = ctx.WithContext(context.WithValue(ctx.Context(), utils.ContextKey("connectionID"), suite.path.EndpointA.ConnectionID))
+				return suite.GetQuicksilverApp(suite.chainA).InterchainstakingKeeper, ctx
+			},
+			wantErr:         false,
+			expectedErr:     nil,
+			expectedChainID: "testchain2",
+		},
+	}
+	for _, tc := range testCase {
+		suite.Run(tc.name, func() {
+			keeper, ctx := tc.setup()
+
+			chainID, err := keeper.GetChainIDFromContext(ctx)
+			if tc.wantErr {
+				suite.Equal(tc.expectedErr, err)
+				return
+			}
+			suite.NoError(err)
+			suite.Equal(tc.expectedChainID, chainID)
+		})
+	}
+}
+
 func (suite *KeeperTestSuite) TestIteratePortConnection() {
 	suite.SetupTest()
 	suite.setupTestZones()