fix tests

cosmos · Dec 2, 2024 · a7ee203 · a7ee203
1 parent 1ed4b64
commit a7ee203
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 4 deletions.
diff --git a/x/accounts/defaults/lockup/continuous_locking_account_test.go b/x/accounts/defaults/lockup/continuous_locking_account_test.go
@@ -111,6 +111,9 @@ func TestContinousAccountUndelegate(t *testing.T) {
 	err = acc.CheckUbdEntriesMature(sdkCtx)
 	require.NoError(t, err)
 
+	_, err = acc.UnbondEntries.Get(sdkCtx, "val_address")
+	require.Error(t, err)
+
 	delLocking, err = acc.DelegatedLocking.Get(ctx, "test")
 	require.NoError(t, err)
 	require.True(t, delLocking.Equal(math.ZeroInt()))

diff --git a/x/accounts/defaults/lockup/delayed_locking_account_test.go b/x/accounts/defaults/lockup/delayed_locking_account_test.go
@@ -110,6 +110,9 @@ func TestDelayedAccountUndelegate(t *testing.T) {
 	err = acc.CheckUbdEntriesMature(sdkCtx)
 	require.NoError(t, err)
 
+	_, err = acc.UnbondEntries.Get(sdkCtx, "val_address")
+	require.Error(t, err)
+
 	delLocking, err = acc.DelegatedLocking.Get(ctx, "test")
 	require.NoError(t, err)
 	require.True(t, delLocking.Equal(math.ZeroInt()))

diff --git a/x/accounts/defaults/lockup/lockup.go b/x/accounts/defaults/lockup/lockup.go
@@ -480,6 +480,7 @@ func (bva *BaseLockup) CheckUbdEntriesMature(ctx context.Context) error {
 
 	currentTime := bva.headerService.HeaderInfo(ctx).Time
 
+	removeKeys := []string{}
 	err = bva.UnbondEntries.Walk(ctx, nil, func(key string, value lockuptypes.UnbondingEntries) (stop bool, err error) {
 		for i := 0; i < len(value.Entries); i++ {
 			entry := value.Entries[i]
@@ -524,10 +525,7 @@ func (bva *BaseLockup) CheckUbdEntriesMature(ctx context.Context) error {
 		}
 
 		if len(value.Entries) == 0 {
-			err = bva.UnbondEntries.Remove(ctx, key)
-			if err != nil {
-				return true, err
-			}
+			removeKeys = append(removeKeys, key)
 		} else {
 			err = bva.UnbondEntries.Set(ctx, key, value)
 			if err != nil {
@@ -537,6 +535,13 @@ func (bva *BaseLockup) CheckUbdEntriesMature(ctx context.Context) error {
 
 		return false, nil
 	})
+
+	for _, key := range removeKeys {
+		err = bva.UnbondEntries.Remove(ctx, key)
+		if err != nil {
+			return err
+		}
+	}
 	return err
 }
 

diff --git a/x/accounts/defaults/lockup/periodic_locking_account_test.go b/x/accounts/defaults/lockup/periodic_locking_account_test.go
@@ -145,6 +145,9 @@ func TestPeriodicAccountUndelegate(t *testing.T) {
 	err = acc.CheckUbdEntriesMature(sdkCtx)
 	require.NoError(t, err)
 
+	_, err = acc.UnbondEntries.Get(sdkCtx, "val_address")
+	require.Error(t, err)
+
 	delLocking, err = acc.DelegatedLocking.Get(ctx, "test")
 	require.NoError(t, err)
 	require.True(t, delLocking.Equal(math.ZeroInt()))

diff --git a/x/accounts/defaults/lockup/permanent_locking_account_test.go b/x/accounts/defaults/lockup/permanent_locking_account_test.go
@@ -87,6 +87,9 @@ func TestPermanentAccountUndelegate(t *testing.T) {
 	err = acc.CheckUbdEntriesMature(sdkCtx)
 	require.NoError(t, err)
 
+	_, err = acc.UnbondEntries.Get(sdkCtx, "val_address")
+	require.Error(t, err)
+
 	delLocking, err = acc.DelegatedLocking.Get(ctx, "test")
 	require.NoError(t, err)
 	require.True(t, delLocking.Equal(math.ZeroInt()))