unbonding was missing setting the locked balance to match what was ju…

…st unbonded

x/gridnode/keeper/keeper.go

@@ -387,6 +387,24 @@ func (k Keeper) GetDelegatedAmount(ctx context.Context, delegator sdk.AccAddress
 	return delegationData.LockedBalance
 }
 
+func (k Keeper) GetPublicKeyForDelegator(ctx context.Context, delegator sdk.AccAddress) (string, error) {
+	storeAdapter := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
+	store := prefix.NewStore(storeAdapter, []byte(types.StoreKey))
+
+	bz := store.Get(k.keyForDelegator(delegator))
+	if bz == nil {
+		return "", errors.New("delegator not found")
+	}
+
+	var delegationData types.DelegationData
+	err := json.Unmarshal(bz, &delegationData)
+	if err != nil {
+		return "", err
+	}
+
+	return delegationData.PublicKey, nil
+}
+
 func (k Keeper) SetDelegatedAmount(ctx context.Context, delegator sdk.AccAddress, amount sdkmath.Int) {
 	storeAdapter := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
 	store := prefix.NewStore(storeAdapter, []byte(types.StoreKey))

x/gridnode/module/abci.go

@@ -56,25 +56,48 @@ func BeginBlocker(goCtx context.Context, k keeper.Keeper) {
 
 			if currentTime.After(timestamp) {
 				fmt.Printf("Processing unbonding for delegator: %s, amount: %d\n", entry.Account, entry.Amount)
+				// Fetch the public key along with the current locked balance
+
 				bankKeeper := k.GetBankKeeper()
 				// Process the unbonding
 				delegatorAddr, err := sdk.AccAddressFromBech32(entry.Account)
 				if err != nil {
 					fmt.Printf("Error processing unbonding for delegator %s: %v\n", entry.Account, err)
 					continue
 				}
+
 				amount := math.NewInt(entry.Amount)
 				coin := sdk.NewCoin("uugd", amount)
 				snd := bankKeeper.SendCoinsFromModuleToAccount(goCtx, types.ModuleName, delegatorAddr, sdk.NewCoins(coin))
 				if snd != nil {
 					fmt.Println("Error sending coins from module to account:", err)
 					continue
 				}
-				// Reduce the delegated amount from the store
+				// Get the current delegated and locked amounts
 				currentDelegatedAmount := k.GetDelegatedAmount(goCtx, delegatorAddr)
+				currentLockedBalance := k.GetLockedBalance(goCtx, delegatorAddr)
+
+				// Calculate new amounts after unbonding
 				newDelegatedAmount := currentDelegatedAmount.Sub(amount)
+				newLockedBalance := currentLockedBalance.Sub(amount)
+
+				// Get the public key for the delegator
+				accountPublicKey, errPk := k.GetPublicKeyForDelegator(goCtx, delegatorAddr)
+				if errPk != nil {
+					fmt.Printf("Error retrieving public key for delegator %s: %v\n", entry.Account, errPk)
+					continue
+				}
+				fmt.Printf("Public key for delegator %s: %s\n", entry.Account, accountPublicKey)
+
+				// Update the stored delegated amount
 				k.SetDelegatedAmount(goCtx, delegatorAddr, newDelegatedAmount)
 
+				// Update the stored locked balance with the public key
+				k.SetLockedBalance(goCtx, delegatorAddr, newLockedBalance, accountPublicKey)
+
+				// Log for confirmation
+				fmt.Printf("Updated balance and delegation for delegator %s\n", entry.Account)
+
 				// Placeholder to call hedgehog
 				fmt.Printf("Placeholder: Notify hedgehog that account %s is unbonding %d tokens.\n", entry.Account, entry.Amount)
 				// TODO: Implement the actual call to hedgehog here