[LUM-863] Entities fix (#67)

* Make sure to back save entities

* Added migration to fix entities

* code format

* Format

x/millions/keeper/keeper_deposit.go

@@ -388,3 +388,14 @@ func (k Keeper) ListDeposits(ctx sdk.Context) (deposits []types.Deposit) {
 	}
 	return
 }
+
+func (k Keeper) UnsafeSetUnpersistedDeposits(ctx sdk.Context) int {
+	i := 0
+	for _, deposit := range k.ListDeposits(ctx) {
+		k.setPoolDeposit(ctx, &deposit)
+		k.setAccountDeposit(ctx, &deposit)
+		i++
+	}
+
+	return i
+}

x/millions/keeper/keeper_withdrawal.go

@@ -416,3 +416,14 @@ func (k Keeper) ListPoolWithdrawals(ctx sdk.Context, poolID uint64) (withdrawals
 	}
 	return
 }
+
+func (k Keeper) UnsafeSetUnpersistedWithdrawals(ctx sdk.Context) int {
+	i := 0
+	for _, withdrawal := range k.ListWithdrawals(ctx) {
+		k.setPoolWithdrawal(ctx, withdrawal)
+		k.setAccountWithdrawal(ctx, withdrawal)
+		i++
+	}
+
+	return i
+}

x/millions/keeper/msg_server_pool.go

@@ -78,10 +78,11 @@ func (k msgServer) restoreICADepositEntities(ctx sdk.Context, poolID uint64) {
 	// Restore deposits ICA locked operations on ICADeposit account
 	deposits := k.Keeper.ListPoolDeposits(ctx, poolID)
 	for _, d := range deposits {
-		if d.State == types.DepositState_IcaDelegate {
+		if d.State == types.DepositState_IcaDelegate || d.State == types.DepositState_IbcTransfer {
 			d.ErrorState = d.State
 			d.State = types.DepositState_Failure
 			k.Keeper.setPoolDeposit(ctx, &d)
+			k.Keeper.setAccountDeposit(ctx, &d)
 		}
 	}
 	// Restore withdrawals ICA locked operations on ICADeposit account
@@ -91,6 +92,7 @@ func (k msgServer) restoreICADepositEntities(ctx sdk.Context, poolID uint64) {
 			w.ErrorState = w.State
 			w.State = types.WithdrawalState_Failure
 			k.Keeper.setPoolWithdrawal(ctx, w)
+			k.Keeper.setAccountWithdrawal(ctx, w)
 		}
 	}
 	// Restore draws ICA locked operations on ICADeposit account

x/millions/migrations/migrator.go

@@ -42,5 +42,13 @@ func (m Migrator) Migrate4To5(ctx sdk.Context) error {
 
 // Migrate5To6 migrates from version 5 to 6
 func (m Migrator) Migrate5To6(ctx sdk.Context) error {
-	return v164.InitializePoolFeeTakers(ctx, m.keeper)
+	if err := v164.SaveEntitiesOnAccountLevel(ctx, m.keeper); err != nil {
+		return err
+	}
+
+	if err := v164.InitializePoolFeeTakers(ctx, m.keeper); err != nil {
+		return err
+	}
+
+	return nil
 }

x/millions/migrations/v164/store.go

@@ -5,9 +5,21 @@ import (
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 
 	millionskeeper "github.com/lum-network/chain/x/millions/keeper"
+
 	millionstypes "github.com/lum-network/chain/x/millions/types"
 )
 
+func SaveEntitiesOnAccountLevel(ctx sdk.Context, k millionskeeper.Keeper) error {
+	ctx.Logger().Info("Saving entities on account level")
+
+	updatedDeposits := k.UnsafeSetUnpersistedDeposits(ctx)
+	updatedWithdrawals := k.UnsafeSetUnpersistedWithdrawals(ctx)
+
+	ctx.Logger().Info("Saving entities on account level: deposits", "count", updatedDeposits)
+	ctx.Logger().Info("Saving entities on account level: withdrawals", "count", updatedWithdrawals)
+	return nil
+}
+
 func InitializePoolFeeTakers(ctx sdk.Context, k millionskeeper.Keeper) error {
 	ctx.Logger().Info("Initializing pool fee takers...")