refactor rebalance to workflow

x/liquidstakeibc/keeper/hooks.go

@@ -117,7 +117,7 @@ func (k *Keeper) AfterEpochEnd(ctx sdk.Context, epochIdentifier string, epochNum
 	}
 
 	if epochIdentifier == liquidstakeibctypes.RedelegationEpochIdentifer {
-		k.Rebalance(ctx, epochNumber)
+		k.RebalanceWorkflow(ctx, epochNumber)
 	}
 
 	return nil
@@ -883,3 +883,32 @@ func (k *Keeper) LSMWorkflow(ctx sdk.Context) {
 		)
 	}
 }
+
+// RebalanceWorkflow tries to make redelegate transactions to host-chain to balance the delegations as per the weights.
+func (k Keeper) RebalanceWorkflow(ctx sdk.Context, epoch int64) {
+
+	hcs := k.GetAllHostChains(ctx)
+	for _, hc := range hcs {
+		// skip unbonding epoch, as we do not want to redelegate tokens that might be going through unbond txn in same epoch.
+		// nothing bad will happen even if we do as long as unbonding txns are triggered before redelegations.
+		if !liquidstakeibctypes.IsUnbondingEpoch(hc.UnbondingFactor, epoch) {
+			k.Logger(ctx).Info("redelegation epoch co-incides with unbonding epoch, skipping it")
+			continue
+		}
+		msgs := k.GenerateRedelegateMsgs(ctx, *hc)
+		// send one msg per ica
+		for _, msg := range msgs {
+			ibcSeq, err := k.GenerateAndExecuteICATx(ctx, hc.ConnectionId, hc.DelegationAccount.Owner, []proto.Message{msg})
+			if err != nil {
+				k.Logger(ctx).Error("Failed to submit ica redelegate txns with", "err:", err)
+				continue
+			}
+			k.SetRedelegationTx(ctx, &liquidstakeibctypes.RedelegateTx{
+				ChainId:       hc.ChainId,
+				IbcSequenceId: ibcSeq,
+				State:         liquidstakeibctypes.RedelegateTx_REDELEGATE_SENT,
+			})
+		}
+	}
+	return
+}

x/liquidstakeibc/keeper/rebalance.go

@@ -19,34 +19,6 @@ type delegation struct {
 	validatorDetails types.Validator
 }
 
-// Rebalance tries to make redelegate transactions to host-chain to balance the delegations as per the weights.
-func (k Keeper) Rebalance(ctx sdk.Context, epoch int64) []proto.Message {
-
-	hcs := k.GetAllHostChains(ctx)
-	for _, hc := range hcs {
-		// skip unbonding epoch, as we do not want to redelegate tokens that might be going through unbond txn in same epoch.
-		// nothing bad will happen even if we do as long as unbonding txns are triggered before redelegations.
-		if !types.IsUnbondingEpoch(hc.UnbondingFactor, epoch) {
-			k.Logger(ctx).Info("redelegation epoch co-incides with unbonding epoch, skipping it")
-			continue
-		}
-		msgs := k.GenerateRedelegateMsgs(ctx, *hc)
-		// send one msg per ica
-		for _, msg := range msgs {
-			ibcSeq, err := k.GenerateAndExecuteICATx(ctx, hc.ConnectionId, hc.DelegationAccount.Owner, []proto.Message{msg})
-			if err != nil {
-				k.Logger(ctx).Error("Failed to submit ica redelegate txns with", "err:", err)
-			}
-			k.SetRedelegationTx(ctx, &types.RedelegateTx{
-				ChainId:       hc.ChainId,
-				IbcSequenceId: ibcSeq,
-				State:         types.RedelegateTx_REDELEGATE_SENT,
-			})
-		}
-	}
-	return nil
-}
-
 func (k Keeper) GenerateRedelegateMsgs(ctx sdk.Context, hc types.HostChain) []proto.Message {
 	var AcceptableDelta = hc.Params.RedelegationAcceptableDelta
 	var MaxRedelegationEntries = hc.Params.MaxEntries

x/liquidstakeibc/keeper/rebalance_test.go

@@ -73,8 +73,8 @@ func (suite *IntegrationTestSuite) TestKeeper_Rebalance() {
 			if got := k.GenerateRedelegateMsgs(suite.ctx, *hc); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("GenerateRedelegateMsgs() = %v, want %v", got, tt.want)
 			}
-			suite.NotPanics(func() { k.Rebalance(suite.ctx, hc.UnbondingFactor) })
-			suite.NotPanics(func() { k.Rebalance(suite.ctx, hc.UnbondingFactor+1) })
+			suite.NotPanics(func() { k.RebalanceWorkflow(suite.ctx, hc.UnbondingFactor) })
+			suite.NotPanics(func() { k.RebalanceWorkflow(suite.ctx, hc.UnbondingFactor+1) })
 		})
 	}
 }