unstake all sanity check test

loadtest/sanitycheck/base_sanity_check.go

@@ -134,10 +134,15 @@ func (t *BaseSanityCheckTest) approveNativeERC20(sender *crypto.ECDSAKey,
 }
 
 // waitForEndOfEpoch waits for the end of the current epoch.
-func (t *BaseSanityCheckTest) waitForEpochEnding() (*types.Header, error) {
+func (t *BaseSanityCheckTest) waitForEpochEnding(fromBlock *uint64) (*types.Header, error) {
 	fmt.Println("Waiting for end of epoch")
 
-	currentBlock, err := t.client.GetBlockByNumber(jsonrpc.LatestBlockNumber, false)
+	rpcBlock := jsonrpc.LatestBlockNumber
+	if fromBlock != nil {
+		rpcBlock = jsonrpc.BlockNumber(*fromBlock)
+	}
+
+	currentBlock, err := t.client.GetBlockByNumber(rpcBlock, false)
 	if err != nil {
 		return nil, err
 	}

loadtest/sanitycheck/sanity_check_register_validator.go

@@ -52,60 +52,67 @@ func (t *RegisterValidatorTest) Run() error {
 	fmt.Println("Running", t.Name())
 	defer fmt.Println("Finished", t.Name())
 
+	_, err := t.runTest()
+
+	return err
+}
+
+// runTest runs the register validator test.
+func (t *RegisterValidatorTest) runTest() (*wallet.Account, error) {
 	fundAmount := ethgo.Ether(2)
 	stakeAmount := ethgo.Ether(1)
 
 	newValidatorAcc, err := wallet.GenerateAccount()
 	if err != nil {
-		return fmt.Errorf("failed to generate new validator key: %w", err)
+		return nil, fmt.Errorf("failed to generate new validator key: %w", err)
 	}
 
 	if err := t.fundAddress(newValidatorAcc.Address(), fundAmount); err != nil {
-		return fmt.Errorf("failed to fund new validator address: %w", err)
+		return nil, fmt.Errorf("failed to fund new validator address: %w", err)
 	}
 
 	bladeAdminKey, err := t.decodePrivateKey(t.config.ValidatorKeys[0])
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	if err := t.whitelistValidators(bladeAdminKey, newValidatorAcc.Address()); err != nil {
-		return fmt.Errorf("failed to whitelist new validator: %w", err)
+		return nil, fmt.Errorf("failed to whitelist new validator: %w", err)
 	}
 
 	if err := t.registerValidator(newValidatorAcc, stakeAmount); err != nil {
-		return fmt.Errorf("failed to register new validator: %w", err)
+		return nil, fmt.Errorf("failed to register new validator: %w", err)
 	}
 
-	epochEndingBlock, err := t.waitForEpochEnding()
+	epochEndingBlock, err := t.waitForEpochEnding(nil)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	extra, err := polybft.GetIbftExtra(epochEndingBlock.ExtraData)
 	if err != nil {
-		return fmt.Errorf("failed to get ibft extra data for epoch ending block. Error: %w", err)
+		return nil, fmt.Errorf("failed to get ibft extra data for epoch ending block. Error: %w", err)
 	}
 
 	fmt.Println("Checking if new validator is added to validator set with its stake")
 
 	if extra.Validators == nil || extra.Validators.IsEmpty() {
-		return fmt.Errorf("validator set delta is empty on an epoch ending block")
+		return nil, fmt.Errorf("validator set delta is empty on an epoch ending block")
 	}
 
 	if !extra.Validators.Added.ContainsAddress(newValidatorAcc.Address()) {
-		return fmt.Errorf("validator %s is not in the added validators", newValidatorAcc.Address())
+		return nil, fmt.Errorf("validator %s is not in the added validators", newValidatorAcc.Address())
 	}
 
 	validatorMetaData := extra.Validators.Added.GetValidatorMetadata(newValidatorAcc.Address())
 	if validatorMetaData.VotingPower.Cmp(stakeAmount) != 0 {
-		return fmt.Errorf("voting power of validator %s is incorrect. Expected: %s, Actual: %s",
+		return nil, fmt.Errorf("voting power of validator %s is incorrect. Expected: %s, Actual: %s",
 			newValidatorAcc.Address(), stakeAmount, validatorMetaData.VotingPower)
 	}
 
 	fmt.Println("Validator", newValidatorAcc.Address(), "is added to the new validator set with correct voting power")
 
-	return nil
+	return newValidatorAcc, nil
 }
 
 // whitelistValidators adds the given validators to the whitelist on StakeManager contract.

loadtest/sanitycheck/sanity_check_runner.go

@@ -117,11 +117,17 @@ func registerTests(cfg *SanityCheckTestConfig,
 		return nil, err
 	}
 
+	unstakeAllTest, err := NewUnstakeAllTest(cfg, testAccountKey, client)
+	if err != nil {
+		return nil, err
+	}
+
 	return []SanityCheckTest{
 		stakeTest,
 		unstakeTest,
 		registerValidatorTest,
 		withdrawRewardsTest,
+		unstakeAllTest,
 	}, nil
 }
 

loadtest/sanitycheck/sanity_check_stake.go

@@ -84,7 +84,7 @@ func (t *StakeTest) Run() error {
 		return fmt.Errorf("stake amount is incorrect. Expected: %s, Actual: %s", expectedStake, currentStake)
 	}
 
-	epochEndingBlock, err := t.waitForEpochEnding()
+	epochEndingBlock, err := t.waitForEpochEnding(nil)
 	if err != nil {
 		return err
 	}

loadtest/sanitycheck/sanity_check_unstake.go

@@ -67,7 +67,7 @@ func (t *UnstakeTest) Run() error {
 
 	fmt.Println("Stake of validator", validatorKey.Address(), "before unstaking:", previousStake)
 
-	if err := t.unstake(validatorKey, amountToUnstake); err != nil {
+	if _, err := t.unstake(validatorKey, amountToUnstake); err != nil {
 		return fmt.Errorf("failed to stake for validator: %s. Error: %w", validatorKey.Address(), err)
 	}
 
@@ -83,7 +83,7 @@ func (t *UnstakeTest) Run() error {
 		return fmt.Errorf("stake amount is incorrect. Expected: %s, Actual: %s", expectedStake, currentStake)
 	}
 
-	epochEndingBlock, err := t.waitForEpochEnding()
+	epochEndingBlock, err := t.waitForEpochEnding(nil)
 	if err != nil {
 		return err
 	}
@@ -115,12 +115,12 @@ func (t *UnstakeTest) Run() error {
 }
 
 // unstake unstakes the given amount for the given validator.
-func (t *UnstakeTest) unstake(validatorKey *crypto.ECDSAKey, amount *big.Int) error {
+func (t *UnstakeTest) unstake(validatorKey *crypto.ECDSAKey, amount *big.Int) (uint64, error) {
 	fmt.Println("Unstaking for validator", validatorKey.Address(), "Amount", amount.String())
 
 	s := time.Now().UTC()
 	defer func() {
-		fmt.Println("Staking for validator", validatorKey.Address(), "took", time.Since(s))
+		fmt.Println("Unstaking for validator", validatorKey.Address(), "took", time.Since(s))
 	}()
 
 	unstakeFn := &contractsapi.UnstakeStakeManagerFn{
@@ -129,7 +129,7 @@ func (t *UnstakeTest) unstake(validatorKey *crypto.ECDSAKey, amount *big.Int) er
 
 	encoded, err := unstakeFn.EncodeAbi()
 	if err != nil {
-		return err
+		return 0, err
 	}
 
 	tx := types.NewTx(types.NewLegacyTx(
@@ -139,12 +139,12 @@ func (t *UnstakeTest) unstake(validatorKey *crypto.ECDSAKey, amount *big.Int) er
 
 	receipt, err := t.txrelayer.SendTransaction(tx, validatorKey)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
 	if receipt.Status == uint64(types.ReceiptFailed) {
-		return fmt.Errorf("unstake transaction failed on block %d", receipt.BlockNumber)
+		return 0, fmt.Errorf("unstake transaction failed on block %d", receipt.BlockNumber)
 	}
 
-	return nil
+	return receipt.BlockNumber, nil
 }

loadtest/sanitycheck/sanity_check_unstake_all.go

@@ -0,0 +1,97 @@
+package sanitycheck
+
+import (
+	"fmt"
+
+	"github.com/0xPolygon/polygon-edge/consensus/polybft"
+	"github.com/0xPolygon/polygon-edge/crypto"
+	"github.com/0xPolygon/polygon-edge/jsonrpc"
+	"github.com/0xPolygon/polygon-edge/types"
+	"github.com/umbracle/ethgo"
+)
+
+// UnstakeAllTest is a test that unstakes all the stake of a validator
+type UnstakeAllTest struct {
+	*UnstakeTest
+	*RegisterValidatorTest
+}
+
+// NewUnstakeAllTest creates a new UnstakeAllTest
+func NewUnstakeAllTest(cfg *SanityCheckTestConfig,
+	testAccountKey *crypto.ECDSAKey, client *jsonrpc.EthClient) (*UnstakeAllTest, error) {
+	unstakeTest, err := NewUnstakeTest(cfg, testAccountKey, client)
+	if err != nil {
+		return nil, err
+	}
+
+	registerValidatorTest, err := NewRegisterValidatorTest(cfg, testAccountKey, client)
+	if err != nil {
+		return nil, err
+	}
+
+	return &UnstakeAllTest{
+		UnstakeTest:           unstakeTest,
+		RegisterValidatorTest: registerValidatorTest,
+	}, nil
+}
+
+// Name returns the name of the unstake all test
+func (t *UnstakeAllTest) Name() string {
+	return "Unstake All Test"
+}
+
+// Name returns the name of the unstake all test
+// It does the following steps:
+// 1. Register a new validator.
+// 2. Unstake all the stake of the validator.
+// 3. Wait for the epoch ending block.
+// 4. Check if the validator is removed from the validator set.
+func (t *UnstakeAllTest) Run() error {
+	printUxSeparator()
+
+	fmt.Println("Running", t.Name())
+	defer fmt.Println("Finished", t.Name())
+
+	validatorAcc, err := t.RegisterValidatorTest.runTest()
+	if err != nil {
+		return err
+	}
+
+	blockNum, err := t.UnstakeTest.unstake(validatorAcc.Ecdsa, ethgo.Ether(1))
+	if err != nil {
+		return err
+	}
+
+	var epochEndingBlock *types.Header
+
+	if blockNum%t.config.EpochSize != 0 {
+		epochEndingBlock, err = t.waitForEpochEnding(&blockNum)
+		if err != nil {
+			return err
+		}
+	} else {
+		epochEndingBlock, err = t.client.GetHeaderByNumber(jsonrpc.BlockNumber(blockNum))
+		if err != nil {
+			return err
+		}
+	}
+
+	extra, err := polybft.GetIbftExtra(epochEndingBlock.ExtraData)
+	if err != nil {
+		return fmt.Errorf("failed to get ibft extra data for epoch ending block. Error: %w", err)
+	}
+
+	fmt.Println("Checking if validator is removed from validator set since it unstaked all")
+
+	if extra.Validators == nil || extra.Validators.IsEmpty() {
+		return fmt.Errorf("validator set delta is empty on an epoch ending block")
+	}
+
+	if len(extra.Validators.Removed) != 1 {
+		return fmt.Errorf("expected 1 validator to be removed from the validator set, got %d", len(extra.Validators.Removed))
+	}
+
+	fmt.Println("Validator", validatorAcc.Address(), "is removed from the validator set")
+
+	return nil
+}

loadtest/sanitycheck/sanity_check_withdraw_rewards.go

@@ -60,7 +60,7 @@ func (t *WithdrawRewardsTest) Run() error {
 	}
 
 	// lets wait for one epoch so that there are some rewards accumulated
-	_, err = t.waitForEpochEnding()
+	_, err = t.waitForEpochEnding(nil)
 	if err != nil {
 		return err
 	}