Fixed forced_undelegate logic; amend the test

execution_engine_testing/tests/src/test/system_contracts/auction/bids.rs

@@ -823,8 +823,7 @@ fn should_forcibly_undelegate_after_setting_validator_limits() {
     builder.forced_undelegate(None, DEFAULT_PROTOCOL_VERSION, DEFAULT_BLOCK_TIME);
 
     let bids = builder.get_bids();
-    // The undelegation itself doesn't remove bids, only process_unbond does.
-    assert_eq!(bids.len(), 3);
+    assert_eq!(bids.len(), 2);
 
     assert!(builder.get_validator_weights(new_era + 1).is_none());
 

storage/src/system/auction.rs

@@ -441,20 +441,32 @@ pub trait Auction:
                         amount,
                         None,
                     )?;
-                    match delegator.decrease_stake(amount, era_end_timestamp_millis) {
-                        Ok(_) => (),
-                        // Work around the case when the locked amounts table has yet to be
-                        // initialized (likely pre-90 day mark).
-                        Err(Error::DelegatorFundsLocked) => continue,
-                        Err(err) => return Err(err),
-                    }
+                    let updated_stake =
+                        match delegator.decrease_stake(amount, era_end_timestamp_millis) {
+                            Ok(updated_stake) => updated_stake,
+                            // Work around the case when the locked amounts table has yet to be
+                            // initialized (likely pre-90 day mark).
+                            Err(Error::DelegatorFundsLocked) => continue,
+                            Err(err) => return Err(err),
+                        };
                     let delegator_bid_addr = BidAddr::new_from_public_keys(
                         validator_public_key,
                         Some(&delegator_public_key),
                     );
 
-                    debug!("pruning delegator bid {}", delegator_bid_addr);
-                    self.prune_bid(delegator_bid_addr);
+                    if updated_stake.is_zero() {
+                        debug!("pruning delegator bid {}", delegator_bid_addr);
+                        self.prune_bid(delegator_bid_addr);
+                    } else {
+                        debug!(
+                            "forced undelegation for {} reducing {} by {} to {}",
+                            delegator_bid_addr, staked_amount, amount, updated_stake
+                        );
+                        self.write_bid(
+                            delegator_bid_addr.into(),
+                            BidKind::Delegator(Box::new(delegator)),
+                        )?;
+                    }
                 }
             }
         }