From 572386349ff3168ca326265dde655ccaa2d79fb9 Mon Sep 17 00:00:00 2001
From: rustdev <placex.com@gmail.com>
Date: Mon, 18 Dec 2023 12:25:35 +0000
Subject: [PATCH] use staking middleware in custom staking module to batch
 delegation.

---
 app/keepers/keepers.go                  |  6 +++---
 custom/staking/keeper/keeper.go         | 22 +++++++++++++---------
 custom/staking/keeper/msg_server.go     |  1 +
 x/stakingmiddleware/client/cli/query.go | 11 +++++++----
 4 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/app/keepers/keepers.go b/app/keepers/keepers.go
index 49d7933ec..f7a155c18 100644
--- a/app/keepers/keepers.go
+++ b/app/keepers/keepers.go
@@ -193,12 +193,12 @@ func (appKeepers *AppKeepers) InitNormalKeepers(
 		appKeepers.AccountKeeper, appKeepers.BankKeeper, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
+	appKeepers.StakingMiddlewareKeeper = stakingmiddleware.NewKeeper(appCodec, appKeepers.keys[stakingmiddlewaretypes.StoreKey], authtypes.NewModuleAddress(govtypes.ModuleName).String())
+
 	appKeepers.CustomStakingKeeper = customstaking.NewKeeper(
-		appCodec /*appKeepers.keys[stakingtypes.StoreKey],*/, *appKeepers.StakingKeeper, appKeepers.AccountKeeper, &appKeepers.MintKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
+		appCodec /*appKeepers.keys[stakingtypes.StoreKey],*/, *appKeepers.StakingKeeper, appKeepers.AccountKeeper, &appKeepers.MintKeeper, &appKeepers.StakingMiddlewareKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
-	appKeepers.StakingMiddlewareKeeper = stakingmiddleware.NewKeeper(appCodec, appKeepers.keys[stakingmiddlewaretypes.StoreKey], authtypes.NewModuleAddress(govtypes.ModuleName).String())
-
 	appKeepers.DistrKeeper = distrkeeper.NewKeeper(
 		appCodec, appKeepers.keys[distrtypes.StoreKey], appKeepers.AccountKeeper, appKeepers.BankKeeper,
 		appKeepers.StakingKeeper, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
diff --git a/custom/staking/keeper/keeper.go b/custom/staking/keeper/keeper.go
index 33c843e4d..7717e2249 100644
--- a/custom/staking/keeper/keeper.go
+++ b/custom/staking/keeper/keeper.go
@@ -6,14 +6,16 @@ import (
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 
 	mintkeeper "github.com/notional-labs/composable/v6/x/mint/keeper"
+	stakingmiddleware "github.com/notional-labs/composable/v6/x/stakingmiddleware/keeper"
 )
 
 type Keeper struct {
 	stakingkeeper.Keeper
-	cdc        codec.BinaryCodec
-	acck       accountkeeper.AccountKeeper
-	mintkeeper *mintkeeper.Keeper
-	authority  string
+	cdc               codec.BinaryCodec
+	acck              accountkeeper.AccountKeeper
+	mintkeeper        *mintkeeper.Keeper
+	stakingmiddleware *stakingmiddleware.Keeper
+	authority         string
 }
 
 // func NewBaseKeeper(
@@ -36,14 +38,16 @@ func NewKeeper(
 	staking stakingkeeper.Keeper,
 	acck accountkeeper.AccountKeeper,
 	mintkeeper *mintkeeper.Keeper,
+	stakingmiddleware *stakingmiddleware.Keeper,
 	authority string,
 ) Keeper {
 	keeper := Keeper{
-		Keeper:     staking,
-		acck:       acck,
-		authority:  authority,
-		mintkeeper: mintkeeper,
-		cdc:        cdc,
+		Keeper:            staking,
+		acck:              acck,
+		authority:         authority,
+		mintkeeper:        mintkeeper,
+		stakingmiddleware: stakingmiddleware,
+		cdc:               cdc,
 	}
 	return keeper
 }
diff --git a/custom/staking/keeper/msg_server.go b/custom/staking/keeper/msg_server.go
index d3ba7396a..22a208cef 100644
--- a/custom/staking/keeper/msg_server.go
+++ b/custom/staking/keeper/msg_server.go
@@ -45,6 +45,7 @@ func (k msgServer) Delegate(goCtx context.Context, msg *types.MsgDelegate) (*typ
 	// 	Shares:           msg.Amount.Amount.ToLegacyDec(),
 	// }
 	k.mintkeeper.SetLastTotalPower(ctx, math.Int{})
+	k.stakingmiddleware.SetLastTotalPower(ctx, math.Int{})
 
 	return &types.MsgDelegateResponse{}, nil
 	// return nil, fmt.Errorf("My custom error: Nikita")
diff --git a/x/stakingmiddleware/client/cli/query.go b/x/stakingmiddleware/client/cli/query.go
index a28e31fae..939b8a5c0 100644
--- a/x/stakingmiddleware/client/cli/query.go
+++ b/x/stakingmiddleware/client/cli/query.go
@@ -21,7 +21,7 @@ func GetQueryCmd() *cobra.Command {
 	}
 
 	mintingQueryCmd.AddCommand(
-		GetCmdQueryParams(),
+		GetCmdQueryPower(),
 	)
 
 	return mintingQueryCmd
@@ -29,16 +29,19 @@ func GetQueryCmd() *cobra.Command {
 
 // GetCmdQueryParams implements a command to return the current minting
 // parameters.
-func GetCmdQueryParams() *cobra.Command {
+func GetCmdQueryPower() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "params",
-		Short: "Query the current minting parameters",
+		Use:   "power",
+		Short: "Query the current power",
 		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			clientCtx, err := client.GetClientQueryContext(cmd)
 			if err != nil {
 				return err
 			}
+
+			return clientCtx.PrintString(fmt.Sprintf("%s\n", "hello world"))
+
 			queryClient := types.NewQueryClient(clientCtx)
 
 			params := &types.QueryPowerRequest{}