From ba38f70e34714af1b2137cb09be06777a4a98580 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Mon, 2 Dec 2024 14:35:41 +0100
Subject: [PATCH 1/2] Rebase and fix test asserts

---
 x/accounts/defaults/lockup/lockup_test.go | 65 ++++++++++++++++++++++-
 1 file changed, 63 insertions(+), 2 deletions(-)

diff --git a/x/accounts/defaults/lockup/lockup_test.go b/x/accounts/defaults/lockup/lockup_test.go
index 0748a17e8caf..5ab674802b7e 100644
--- a/x/accounts/defaults/lockup/lockup_test.go
+++ b/x/accounts/defaults/lockup/lockup_test.go
@@ -2,14 +2,15 @@ package lockup
 
 import (
 	"context"
+	"cosmossdk.io/core/header"
+	"cosmossdk.io/log"
+	"github.com/stretchr/testify/assert"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 
-	"cosmossdk.io/core/header"
 	"cosmossdk.io/core/store"
-	"cosmossdk.io/log"
 	"cosmossdk.io/math"
 	lockuptypes "cosmossdk.io/x/accounts/defaults/lockup/v1"
 
@@ -281,4 +282,64 @@ func TestQueryLockupAccountBaseInfo(t *testing.T) {
 	require.Equal(t, res.OriginalLocking.AmountOf("test"), math.NewInt(10))
 	require.Equal(t, res.Owner, "owner")
 	require.NoError(t, err)
+
+}
+
+func TestWithdrawUnlockedCoins(t *testing.T) {
+	const (
+		liquid         = 6
+		locked         = 4
+		initialBalance = liquid + locked // note: 10 is setup in the bank mock
+		staked         = 7
+		unstaked       = staked
+		denom          = "test"
+	)
+	ctx, ss := newMockContext(t)
+	now := time.Now()
+	sdkCtx := sdk.NewContext(nil, true, log.NewNopLogger()).WithContext(ctx).WithHeaderInfo(header.Info{
+		Time: now,
+	})
+
+	baseLockup := setup(t, ctx, ss)
+	require.NoError(t, baseLockup.OriginalLocking.Set(ctx, denom, math.NewInt(locked)))
+
+	foreverLocked := func(ctx context.Context, time time.Time, denoms ...string) (sdk.Coins, error) {
+		return sdk.NewCoins(sdk.NewCoin(denom, math.NewInt(locked))), nil
+	}
+
+	// delegate
+	_, err := baseLockup.Delegate(sdkCtx, &lockuptypes.MsgDelegate{
+		Sender:           "owner",
+		ValidatorAddress: "validator",
+		Amount:           sdk.NewCoin(denom, math.NewInt(staked)),
+	}, foreverLocked)
+	require.NoError(t, err)
+	// balance should be 10 - 7 = 3
+
+	_, err = baseLockup.Undelegate(sdkCtx, &lockuptypes.MsgUndelegate{
+		Sender:           "owner",
+		ValidatorAddress: "validator",
+		Amount:           sdk.NewCoin(denom, math.NewInt(unstaked)),
+	})
+	require.NoError(t, err)
+	// when unbonding time has expired
+	sdkCtx = sdkCtx.WithHeaderInfo(header.Info{Time: now.Add(21*24*time.Hour + time.Minute)})
+
+	// balance should be 3 + 7 = 10 now
+	_, err = baseLockup.UnbondEntries.Get(sdkCtx, "validator")
+	require.NoError(t, err)
+	amount, err := baseLockup.DelegatedLocking.Get(sdkCtx, denom)
+	require.NoError(t, err)
+	assert.Equal(t, amount.String(), math.NewInt(locked).String())
+
+	// when WithdrawUnlockedCoins is called,
+	// only the liquid amount should be returned
+	res, err := baseLockup.WithdrawUnlockedCoins(sdkCtx, &lockuptypes.MsgWithdraw{
+		Withdrawer: "owner",
+		ToAddress:  "my-receiver",
+		Denoms:     []string{denom},
+	}, foreverLocked)
+	// 4 stake is still locked. nothing to return
+	require.Equal(t, math.NewInt(0).String(), res.AmountReceived.AmountOf(denom).String())
+	require.NoError(t, err)
 }

From c690a52d771515c443232839eb3760e524a45ab6 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Tue, 3 Dec 2024 12:18:20 +0100
Subject: [PATCH 2/2] Fix mock response

---
 x/accounts/defaults/lockup/utils_test.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/x/accounts/defaults/lockup/utils_test.go b/x/accounts/defaults/lockup/utils_test.go
index 684e43b79bd3..311a3d1d5518 100644
--- a/x/accounts/defaults/lockup/utils_test.go
+++ b/x/accounts/defaults/lockup/utils_test.go
@@ -78,7 +78,8 @@ func newMockContext(t *testing.T) (context.Context, store.KVStoreService) {
 				return &stakingtypes.MsgDelegateResponse{}, nil
 			case "/cosmos.staking.v1beta1.MsgUndelegate":
 				return &stakingtypes.MsgUndelegateResponse{
-					Amount: sdk.NewCoin("test", math.NewInt(1)),
+					CompletionTime: time.Now().Add(stakingtypes.DefaultUnbondingTime),
+					Amount:         msg.(*stakingtypes.MsgUndelegate).Amount,
 				}, nil
 			case "/cosmos.distribution.v1beta1.MsgWithdrawDelegatorReward":
 				return &distrtypes.MsgWithdrawDelegatorRewardResponse{}, nil