From 110199b68add9ac5b0964b516cd11cf39f5ce3f0 Mon Sep 17 00:00:00 2001
From: Kevin <105028215+movekevin@users.noreply.github.com>
Date: Mon, 26 Sep 2022 09:36:17 -0700
Subject: [PATCH] [Tests only] Add an e2e test for staking_contract flow
 (#4537)

---
 aptos-move/e2e-move-tests/src/tests/stake.rs | 87 +++++++++++++++++++-
 aptos-move/e2e-tests/src/account.rs          |  2 +-
 2 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/aptos-move/e2e-move-tests/src/tests/stake.rs b/aptos-move/e2e-move-tests/src/tests/stake.rs
index 9d896322d1ffc..cbc15d3158f17 100644
--- a/aptos-move/e2e-move-tests/src/tests/stake.rs
+++ b/aptos-move/e2e-move-tests/src/tests/stake.rs
@@ -6,7 +6,7 @@ use crate::{
     initialize_staking, join_validator_set, leave_validator_set, rotate_consensus_key,
     setup_staking, unlock_stake, withdraw_stake, MoveHarness,
 };
-use aptos_types::account_address::AccountAddress;
+use aptos_types::account_address::{default_stake_pool_address, AccountAddress};
 use cached_packages::aptos_stdlib;
 use move_deps::move_core_types::language_storage::CORE_CODE_ADDRESS;
 
@@ -105,7 +105,6 @@ fn test_staking_mainnet() {
     assert_success!(setup_staking(&mut harness, &validator, 100_000_000_000_000));
     harness.new_epoch();
     let validator_set = get_validator_set(&harness);
-    println!("validator_set {:?}", validator_set);
     assert_eq!(validator_set.active_validators.len(), 2);
 }
 
@@ -219,3 +218,87 @@ fn test_staking_rewards_pending_inactive() {
         stake_amount + 570
     );
 }
+
+#[test]
+fn test_staking_contract() {
+    let mut harness = MoveHarness::new();
+    let staker = harness.new_account_at(AccountAddress::from_hex_literal("0x11").unwrap());
+    let operator_1 = harness.new_account_at(AccountAddress::from_hex_literal("0x21").unwrap());
+    let operator_2 = harness.new_account_at(AccountAddress::from_hex_literal("0x22").unwrap());
+    let amount = 25_000_000;
+    let staker_address = *staker.address();
+    let operator_1_address = *operator_1.address();
+    let operator_2_address = *operator_2.address();
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_create_staking_contract(
+            operator_1_address,
+            operator_1_address,
+            amount,
+            10,
+            vec![],
+        )
+    ));
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_add_stake(operator_1_address, amount)
+    ));
+
+    // Join validator set.
+    let pool_address = default_stake_pool_address(staker_address, operator_1_address);
+    assert_success!(rotate_consensus_key(
+        &mut harness,
+        &operator_1,
+        pool_address
+    ));
+    assert_success!(join_validator_set(&mut harness, &operator_1, pool_address));
+    harness.new_epoch();
+    let validator_set = get_validator_set(&harness);
+    assert_eq!(
+        validator_set.active_validators[1].account_address,
+        pool_address,
+    );
+
+    // Operator requests commissions.
+    harness.new_block_with_metadata(pool_address, vec![]);
+    harness.new_epoch();
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_request_commission(staker_address, operator_1_address)
+    ));
+
+    // Wait until stake is unlocked.
+    harness.fast_forward(7200);
+    harness.new_epoch();
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_distribute(staker_address, operator_1_address)
+    ));
+
+    // Staker unlocks some stake.
+    harness.new_block_with_metadata(pool_address, vec![]);
+    harness.new_epoch();
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_unlock_stake(operator_1_address, amount)
+    ));
+
+    // Wait until stake is unlocked.
+    harness.fast_forward(7200);
+    harness.new_epoch();
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_distribute(staker_address, operator_1_address)
+    ));
+
+    // Switch operators.
+    assert_success!(harness.run_transaction_payload(
+        &staker,
+        aptos_stdlib::staking_contract_switch_operator_with_same_commission(
+            operator_1_address,
+            operator_2_address,
+        )
+    ));
+    // New operator leaves validator set.
+    leave_validator_set(&mut harness, &operator_2, pool_address);
+}
diff --git a/aptos-move/e2e-tests/src/account.rs b/aptos-move/e2e-tests/src/account.rs
index 97e707598a914..4838847b367e8 100644
--- a/aptos-move/e2e-tests/src/account.rs
+++ b/aptos-move/e2e-tests/src/account.rs
@@ -26,7 +26,7 @@ use move_deps::move_core_types::{
 use vm_genesis::GENESIS_KEYPAIR;
 
 // TTL is 86400s. Initial time was set to 0.
-pub const DEFAULT_EXPIRATION_TIME: u64 = 40_000;
+pub const DEFAULT_EXPIRATION_TIME: u64 = 4_000_000;
 
 /// Details about a Aptos account.
 ///