From d5ac6b883b1416a2fb56777fdc31fc8cdf9740c8 Mon Sep 17 00:00:00 2001 From: Jon Cinque Date: Mon, 13 Dec 2021 13:17:07 +0100 Subject: [PATCH] stake-pool-py: Add simple rebalance bot (#2630) * stake-pool-py: Add simple rebalance bot * Fixup test * Refactor flaky tests * Create vote earlier in test * Duplicate create_vote call causes stall * Wait more aggressively --- ci/py-test-stake-pool.sh | 1 + stake-pool/program/tests/deposit.rs | 179 ---------------- stake-pool/program/tests/deposit_authority.rs | 200 ++++++++++++++++++ stake-pool/program/tests/initialize.rs | 20 -- stake-pool/py/README.md | 66 +++++- stake-pool/py/bot/__init__.py | 0 stake-pool/py/bot/rebalance.py | 135 ++++++++++++ stake-pool/py/stake_pool/instructions.py | 4 +- stake-pool/py/tests/conftest.py | 27 ++- stake-pool/py/tests/test_a_time_sensitive.py | 10 +- stake-pool/py/tests/test_bot_rebalance.py | 79 +++++++ .../py/tests/test_deposit_withdraw_stake.py | 43 +++- 12 files changed, 547 insertions(+), 217 deletions(-) create mode 100644 stake-pool/program/tests/deposit_authority.rs create mode 100644 stake-pool/py/bot/__init__.py create mode 100644 stake-pool/py/bot/rebalance.py create mode 100644 stake-pool/py/tests/test_bot_rebalance.py diff --git a/ci/py-test-stake-pool.sh b/ci/py-test-stake-pool.sh index 34074a7efbd..3329ebd8abb 100755 --- a/ci/py-test-stake-pool.sh +++ b/ci/py-test-stake-pool.sh @@ -9,6 +9,7 @@ python3 -m venv venv source ./venv/bin/activate pip3 install -r requirements.txt check_dirs=( + "bot" "spl_token" "stake" "stake_pool" diff --git a/stake-pool/program/tests/deposit.rs b/stake-pool/program/tests/deposit.rs index 0b7eda1e7ca..b2eb5f5e3e0 100644 --- a/stake-pool/program/tests/deposit.rs +++ b/stake-pool/program/tests/deposit.rs @@ -793,185 +793,6 @@ async fn fail_with_uninitialized_validator_list() {} // TODO #[tokio::test] async fn fail_with_out_of_dated_pool_balances() {} // TODO -#[tokio::test] -async fn success_with_deposit_authority() { - let (mut banks_client, payer, recent_blockhash) = program_test().start().await; - let stake_deposit_authority = Keypair::new(); - let stake_pool_accounts = - StakePoolAccounts::new_with_deposit_authority(stake_deposit_authority); - stake_pool_accounts - .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1) - .await - .unwrap(); - - let validator_stake_account = simple_add_validator_to_pool( - &mut banks_client, - &payer, - &recent_blockhash, - &stake_pool_accounts, - ) - .await; - - let user = Keypair::new(); - let user_stake = Keypair::new(); - let lockup = stake::state::Lockup::default(); - let authorized = stake::state::Authorized { - staker: user.pubkey(), - withdrawer: user.pubkey(), - }; - let _stake_lamports = create_independent_stake_account( - &mut banks_client, - &payer, - &recent_blockhash, - &user_stake, - &authorized, - &lockup, - TEST_STAKE_AMOUNT, - ) - .await; - - create_vote( - &mut banks_client, - &payer, - &recent_blockhash, - &validator_stake_account.validator, - &validator_stake_account.vote, - ) - .await; - delegate_stake_account( - &mut banks_client, - &payer, - &recent_blockhash, - &user_stake.pubkey(), - &user, - &validator_stake_account.vote.pubkey(), - ) - .await; - - // make pool token account - let user_pool_account = Keypair::new(); - create_token_account( - &mut banks_client, - &payer, - &recent_blockhash, - &user_pool_account, - &stake_pool_accounts.pool_mint.pubkey(), - &user.pubkey(), - ) - .await - .unwrap(); - - let error = stake_pool_accounts - .deposit_stake( - &mut banks_client, - &payer, - &recent_blockhash, - &user_stake.pubkey(), - &user_pool_account.pubkey(), - &validator_stake_account.stake_account, - &user, - ) - .await; - assert!(error.is_none()); -} - -#[tokio::test] -async fn fail_without_deposit_authority_signature() { - let (mut banks_client, payer, recent_blockhash) = program_test().start().await; - let stake_deposit_authority = Keypair::new(); - let mut stake_pool_accounts = - StakePoolAccounts::new_with_deposit_authority(stake_deposit_authority); - stake_pool_accounts - .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1) - .await - .unwrap(); - - let validator_stake_account = simple_add_validator_to_pool( - &mut banks_client, - &payer, - &recent_blockhash, - &stake_pool_accounts, - ) - .await; - - let user = Keypair::new(); - let user_stake = Keypair::new(); - let lockup = stake::state::Lockup::default(); - let authorized = stake::state::Authorized { - staker: user.pubkey(), - withdrawer: user.pubkey(), - }; - let _stake_lamports = create_independent_stake_account( - &mut banks_client, - &payer, - &recent_blockhash, - &user_stake, - &authorized, - &lockup, - TEST_STAKE_AMOUNT, - ) - .await; - - create_vote( - &mut banks_client, - &payer, - &recent_blockhash, - &validator_stake_account.validator, - &validator_stake_account.vote, - ) - .await; - delegate_stake_account( - &mut banks_client, - &payer, - &recent_blockhash, - &user_stake.pubkey(), - &user, - &validator_stake_account.vote.pubkey(), - ) - .await; - - // make pool token account - let user_pool_account = Keypair::new(); - create_token_account( - &mut banks_client, - &payer, - &recent_blockhash, - &user_pool_account, - &stake_pool_accounts.pool_mint.pubkey(), - &user.pubkey(), - ) - .await - .unwrap(); - - let wrong_depositor = Keypair::new(); - stake_pool_accounts.stake_deposit_authority = wrong_depositor.pubkey(); - stake_pool_accounts.stake_deposit_authority_keypair = Some(wrong_depositor); - - let error = stake_pool_accounts - .deposit_stake( - &mut banks_client, - &payer, - &recent_blockhash, - &user_stake.pubkey(), - &user_pool_account.pubkey(), - &validator_stake_account.stake_account, - &user, - ) - .await - .unwrap() - .unwrap(); - - match error { - TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => { - assert_eq!( - error_index, - StakePoolError::InvalidStakeDepositAuthority as u32 - ); - } - _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"), - } -} - #[tokio::test] async fn success_with_preferred_deposit() { let ( diff --git a/stake-pool/program/tests/deposit_authority.rs b/stake-pool/program/tests/deposit_authority.rs new file mode 100644 index 00000000000..ace4b5d1fb4 --- /dev/null +++ b/stake-pool/program/tests/deposit_authority.rs @@ -0,0 +1,200 @@ +#![cfg(feature = "test-bpf")] + +mod helpers; + +use { + helpers::*, + solana_program::{instruction::InstructionError, stake}, + solana_program_test::*, + solana_sdk::{ + borsh::try_from_slice_unchecked, + signature::{Keypair, Signer}, + transaction::TransactionError, + }, + spl_stake_pool::{error::StakePoolError, state::StakePool}, +}; + +#[tokio::test] +async fn success_initialize() { + let (mut banks_client, payer, recent_blockhash) = program_test().start().await; + let deposit_authority = Keypair::new(); + let stake_pool_accounts = StakePoolAccounts::new_with_deposit_authority(deposit_authority); + let deposit_authority = stake_pool_accounts.stake_deposit_authority; + stake_pool_accounts + .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1) + .await + .unwrap(); + + // Stake pool now exists + let stake_pool_account = + get_account(&mut banks_client, &stake_pool_accounts.stake_pool.pubkey()).await; + let stake_pool = + try_from_slice_unchecked::(stake_pool_account.data.as_slice()).unwrap(); + assert_eq!(stake_pool.stake_deposit_authority, deposit_authority); + assert_eq!(stake_pool.sol_deposit_authority.unwrap(), deposit_authority); +} + +#[tokio::test] +async fn success_deposit() { + let (mut banks_client, payer, recent_blockhash) = program_test().start().await; + let stake_deposit_authority = Keypair::new(); + let stake_pool_accounts = + StakePoolAccounts::new_with_deposit_authority(stake_deposit_authority); + stake_pool_accounts + .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1) + .await + .unwrap(); + + let validator_stake_account = simple_add_validator_to_pool( + &mut banks_client, + &payer, + &recent_blockhash, + &stake_pool_accounts, + ) + .await; + + let user = Keypair::new(); + let user_stake = Keypair::new(); + let lockup = stake::state::Lockup::default(); + let authorized = stake::state::Authorized { + staker: user.pubkey(), + withdrawer: user.pubkey(), + }; + + let _stake_lamports = create_independent_stake_account( + &mut banks_client, + &payer, + &recent_blockhash, + &user_stake, + &authorized, + &lockup, + TEST_STAKE_AMOUNT, + ) + .await; + + delegate_stake_account( + &mut banks_client, + &payer, + &recent_blockhash, + &user_stake.pubkey(), + &user, + &validator_stake_account.vote.pubkey(), + ) + .await; + + // make pool token account + let user_pool_account = Keypair::new(); + create_token_account( + &mut banks_client, + &payer, + &recent_blockhash, + &user_pool_account, + &stake_pool_accounts.pool_mint.pubkey(), + &user.pubkey(), + ) + .await + .unwrap(); + + let error = stake_pool_accounts + .deposit_stake( + &mut banks_client, + &payer, + &recent_blockhash, + &user_stake.pubkey(), + &user_pool_account.pubkey(), + &validator_stake_account.stake_account, + &user, + ) + .await; + assert!(error.is_none()); +} + +#[tokio::test] +async fn fail_deposit_without_authority_signature() { + let (mut banks_client, payer, recent_blockhash) = program_test().start().await; + let stake_deposit_authority = Keypair::new(); + let mut stake_pool_accounts = + StakePoolAccounts::new_with_deposit_authority(stake_deposit_authority); + stake_pool_accounts + .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1) + .await + .unwrap(); + + let validator_stake_account = simple_add_validator_to_pool( + &mut banks_client, + &payer, + &recent_blockhash, + &stake_pool_accounts, + ) + .await; + + let user = Keypair::new(); + let user_stake = Keypair::new(); + let lockup = stake::state::Lockup::default(); + let authorized = stake::state::Authorized { + staker: user.pubkey(), + withdrawer: user.pubkey(), + }; + + let _stake_lamports = create_independent_stake_account( + &mut banks_client, + &payer, + &recent_blockhash, + &user_stake, + &authorized, + &lockup, + TEST_STAKE_AMOUNT, + ) + .await; + + delegate_stake_account( + &mut banks_client, + &payer, + &recent_blockhash, + &user_stake.pubkey(), + &user, + &validator_stake_account.vote.pubkey(), + ) + .await; + + // make pool token account + let user_pool_account = Keypair::new(); + create_token_account( + &mut banks_client, + &payer, + &recent_blockhash, + &user_pool_account, + &stake_pool_accounts.pool_mint.pubkey(), + &user.pubkey(), + ) + .await + .unwrap(); + + let wrong_depositor = Keypair::new(); + stake_pool_accounts.stake_deposit_authority = wrong_depositor.pubkey(); + stake_pool_accounts.stake_deposit_authority_keypair = Some(wrong_depositor); + + let error = stake_pool_accounts + .deposit_stake( + &mut banks_client, + &payer, + &recent_blockhash, + &user_stake.pubkey(), + &user_pool_account.pubkey(), + &validator_stake_account.stake_account, + &user, + ) + .await + .unwrap() + .unwrap(); + + match error { + TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => { + assert_eq!( + error_index, + StakePoolError::InvalidStakeDepositAuthority as u32 + ); + } + _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"), + } +} diff --git a/stake-pool/program/tests/initialize.rs b/stake-pool/program/tests/initialize.rs index 6fd44848c92..dd260923567 100644 --- a/stake-pool/program/tests/initialize.rs +++ b/stake-pool/program/tests/initialize.rs @@ -1273,23 +1273,3 @@ async fn fail_with_bad_reserve() { ); } } - -#[tokio::test] -async fn success_with_deposit_authority() { - let (mut banks_client, payer, recent_blockhash) = program_test().start().await; - let deposit_authority = Keypair::new(); - let stake_pool_accounts = StakePoolAccounts::new_with_deposit_authority(deposit_authority); - let deposit_authority = stake_pool_accounts.stake_deposit_authority; - stake_pool_accounts - .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1) - .await - .unwrap(); - - // Stake pool now exists - let stake_pool_account = - get_account(&mut banks_client, &stake_pool_accounts.stake_pool.pubkey()).await; - let stake_pool = - try_from_slice_unchecked::(stake_pool_account.data.as_slice()).unwrap(); - assert_eq!(stake_pool.stake_deposit_authority, deposit_authority); - assert_eq!(stake_pool.sol_deposit_authority.unwrap(), deposit_authority); -} diff --git a/stake-pool/py/README.md b/stake-pool/py/README.md index 97e3605c699..49866fb6d23 100644 --- a/stake-pool/py/README.md +++ b/stake-pool/py/README.md @@ -1,3 +1,67 @@ # Stake-Pool Python Bindings -WIP Python bindings to interact with the stake pool program. +Preliminary Python bindings to interact with the stake pool program, enabling +simple stake delegation bots. + +## To do + +* More reference bot implementations +* Add bindings for all stake pool instructions, see `TODO`s in `stake_pool/instructions.py` +* Finish bindings for vote and stake program +* Upstream vote and stake program bindings to https://github.com/michaelhly/solana-py + +## Development + +### Environment Setup + +1. Ensure that Python 3 is installed with `venv`: https://www.python.org/downloads/ +2. (Optional, but highly recommended) Setup and activate a virtual environment: + +``` +$ python3 -m venv venv +$ source venv/bin/activate +``` + +3. Install requirements + +``` +$ pip install -r requirements.txt +``` + +4. Install the Solana tool suite: https://docs.solana.com/cli/install-solana-cli-tools + +### Test + +Testing through `pytest`: + +``` +$ python3 -m pytest +``` + +Note: the tests all run against a `solana-test-validator` with short epochs of 64 +slots (25.6 seconds exactly). Some tests wait for epoch changes, so they take +time, roughly 90 seconds total at the time of this writing. + +### Formatting + +``` +$ flake8 bot spl_token stake stake_pool system tests vote +``` + +### Type Checker + +``` +$ mypy bot stake stake_pool tests vote spl_token system +``` + +## Delegation Bots + +The `./bot` directory contains sample stake pool delegation bot implementations: + +* `rebalance`: simple bot to make the amount delegated to each validator +uniform, while also maintaining some SOL in the reserve if desired. Can be run +with the stake pool address, staker keypair, and SOL to leave in the reserve: + +``` +$ python3 bot/rebalance.py Zg5YBPAk8RqBR9kaLLSoN5C8Uv7nErBz1WC63HTsCPR staker.json 10.5 +``` diff --git a/stake-pool/py/bot/__init__.py b/stake-pool/py/bot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/stake-pool/py/bot/rebalance.py b/stake-pool/py/bot/rebalance.py new file mode 100644 index 00000000000..37e25261d99 --- /dev/null +++ b/stake-pool/py/bot/rebalance.py @@ -0,0 +1,135 @@ +import argparse +import asyncio +import json + +from solana.keypair import Keypair +from solana.publickey import PublicKey +from solana.rpc.async_api import AsyncClient +from solana.rpc.commitment import Confirmed + +from stake.constants import STAKE_LEN +from stake_pool.actions import decrease_validator_stake, increase_validator_stake, update_stake_pool +from stake_pool.state import StakePool, ValidatorList + + +LAMPORTS_PER_SOL: int = 1_000_000_000 +MINIMUM_INCREASE_LAMPORTS: int = LAMPORTS_PER_SOL // 100 + + +async def get_client(endpoint: str) -> AsyncClient: + print(f'Connecting to network at {endpoint}') + async_client = AsyncClient(endpoint=endpoint, commitment=Confirmed) + total_attempts = 10 + current_attempt = 0 + while not await async_client.is_connected(): + if current_attempt == total_attempts: + raise Exception("Could not connect to test validator") + else: + current_attempt += 1 + await asyncio.sleep(1) + return async_client + + +async def rebalance(endpoint: str, stake_pool_address: PublicKey, staker: Keypair, retained_reserve_amount: float): + async_client = await get_client(endpoint) + + resp = await async_client.get_epoch_info(commitment=Confirmed) + epoch = resp['result']['epoch'] + resp = await async_client.get_account_info(stake_pool_address, commitment=Confirmed) + data = resp['result']['value']['data'] + stake_pool = StakePool.decode(data[0], data[1]) + + print(f'Stake pool last update epoch {stake_pool.last_update_epoch}, current epoch {epoch}') + if stake_pool.last_update_epoch != epoch: + print('Updating stake pool') + await update_stake_pool(async_client, staker, stake_pool_address) + resp = await async_client.get_account_info(stake_pool_address, commitment=Confirmed) + data = resp['result']['value']['data'] + stake_pool = StakePool.decode(data[0], data[1]) + + resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN) + stake_rent_exemption = resp['result'] + retained_reserve_lamports = int(retained_reserve_amount * LAMPORTS_PER_SOL) + + resp = await async_client.get_account_info(stake_pool.validator_list, commitment=Confirmed) + data = resp['result']['value']['data'] + validator_list = ValidatorList.decode(data[0], data[1]) + + print('Stake pool stats:') + print(f'* {stake_pool.total_lamports} total lamports') + num_validators = len(validator_list.validators) + print(f'* {num_validators} validators') + print(f'* Retaining {retained_reserve_lamports} lamports in the reserve') + lamports_per_validator = (stake_pool.total_lamports - retained_reserve_lamports) // num_validators + num_increases = sum([ + 1 for validator in validator_list.validators + if validator.transient_stake_lamports == 0 and validator.active_stake_lamports < lamports_per_validator + ]) + total_usable_lamports = stake_pool.total_lamports - retained_reserve_lamports - num_increases * stake_rent_exemption + lamports_per_validator = total_usable_lamports // num_validators + print(f'* {lamports_per_validator} lamports desired per validator') + + futures = [] + for validator in validator_list.validators: + if validator.transient_stake_lamports != 0: + print(f'Skipping {validator.vote_account_address}: {validator.transient_stake_lamports} transient lamports') + else: + if validator.active_stake_lamports > lamports_per_validator: + lamports_to_decrease = validator.active_stake_lamports - lamports_per_validator + if lamports_to_decrease <= stake_rent_exemption: + print(f'Skipping decrease on {validator.vote_account_address}, \ +currently at {validator.active_stake_lamports} lamports, \ +decrease of {lamports_to_decrease} below the rent exmption') + else: + futures.append(decrease_validator_stake( + async_client, staker, staker, stake_pool_address, + validator.vote_account_address, lamports_to_decrease + )) + elif validator.active_stake_lamports < lamports_per_validator: + lamports_to_increase = lamports_per_validator - validator.active_stake_lamports + if lamports_to_increase < MINIMUM_INCREASE_LAMPORTS: + print(f'Skipping increase on {validator.vote_account_address}, \ +currently at {validator.active_stake_lamports} lamports, \ +increase of {lamports_to_increase} less than the minimum of {MINIMUM_INCREASE_LAMPORTS}') + else: + futures.append(increase_validator_stake( + async_client, staker, staker, stake_pool_address, + validator.vote_account_address, lamports_to_increase + )) + else: + print(f'{validator.vote_account_address}: already at {lamports_per_validator}') + + print('Executing strategy') + await asyncio.gather(*futures) + print('Done') + await async_client.close() + + +def keypair_from_file(keyfile_name: str) -> Keypair: + with open(keyfile_name, 'r') as keyfile: + data = keyfile.read() + int_list = json.loads(data) + bytes_list = [value.to_bytes(1, 'little') for value in int_list] + return Keypair.from_secret_key(b''.join(bytes_list)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Rebalance stake evenly between all the validators in a stake pool.') + parser.add_argument('stake_pool', metavar='STAKE_POOL_ADDRESS', type=str, + help='Stake pool to rebalance, given by a public key in base-58,\ + e.g. Zg5YBPAk8RqBR9kaLLSoN5C8Uv7nErBz1WC63HTsCPR') + parser.add_argument('staker', metavar='STAKER_KEYPAIR', type=str, + help='Staker for the stake pool, given by a keypair file, e.g. staker.json') + parser.add_argument('reserve_amount', metavar='RESERVE_AMOUNT', type=float, + help='Amount of SOL to keep in the reserve, e.g. 10.5') + parser.add_argument('--endpoint', metavar='ENDPOINT_URL', type=str, + default='https://api.mainnet-beta.solana.com', + help='RPC endpoint to use, e.g. https://api.mainnet-beta.solana.com') + + args = parser.parse_args() + stake_pool = PublicKey(args.stake_pool) + staker = keypair_from_file(args.staker) + print(f'Rebalancing stake pool {stake_pool}') + print(f'Staker public key: {staker.public_key}') + print(f'Amount to leave in the reserve: {args.reserve_amount} SOL') + asyncio.run(rebalance(args.endpoint, stake_pool, staker, args.reserve_amount)) diff --git a/stake-pool/py/stake_pool/instructions.py b/stake-pool/py/stake_pool/instructions.py index e0f17abf6ee..8f3f1ccbf1c 100644 --- a/stake-pool/py/stake_pool/instructions.py +++ b/stake-pool/py/stake_pool/instructions.py @@ -505,9 +505,9 @@ class InstructionType(IntEnum): InstructionType.CLEANUP_REMOVED_VALIDATOR_ENTRIES: Pass, InstructionType.DEPOSIT_STAKE: Pass, InstructionType.WITHDRAW_STAKE: AMOUNT_LAYOUT, - InstructionType.SET_MANAGER: Pass, + InstructionType.SET_MANAGER: Pass, # TODO InstructionType.SET_FEE: Pass, # TODO - InstructionType.SET_STAKER: Pass, + InstructionType.SET_STAKER: Pass, # TODO InstructionType.DEPOSIT_SOL: AMOUNT_LAYOUT, InstructionType.SET_FUNDING_AUTHORITY: Pass, # TODO InstructionType.WITHDRAW_SOL: AMOUNT_LAYOUT, diff --git a/stake-pool/py/tests/conftest.py b/stake-pool/py/tests/conftest.py index 32fa07975d4..2c972f19d06 100644 --- a/stake-pool/py/tests/conftest.py +++ b/stake-pool/py/tests/conftest.py @@ -17,6 +17,8 @@ from stake_pool.actions import create_all, add_validator_to_pool from stake_pool.state import Fee +NUM_SLOTS_PER_EPOCH: int = 32 + @pytest.fixture(scope="session") def solana_test_validator(): @@ -28,7 +30,7 @@ def solana_test_validator(): "--reset", "--quiet", "--bpf-program", "SPoo1Ku8WFXoNDMHPsrGSTSG1Y47rzgn41SLUNakuHy", f"{old_cwd}/../../target/deploy/spl_stake_pool.so", - "--slots-per-epoch", "64", + "--slots-per-epoch", str(NUM_SLOTS_PER_EPOCH), ],) yield validator.kill() @@ -93,3 +95,26 @@ def payer(event_loop, async_client) -> Keypair: airdrop_lamports = 10_000_000_000 event_loop.run_until_complete(airdrop(async_client, payer.public_key, airdrop_lamports)) return payer + + +class Waiter: + @staticmethod + async def wait_for_next_epoch(async_client: AsyncClient): + resp = await async_client.get_epoch_info(commitment=Confirmed) + current_epoch = resp['result']['epoch'] + next_epoch = current_epoch + while current_epoch == next_epoch: + await asyncio.sleep(1.0) + resp = await async_client.get_epoch_info(commitment=Confirmed) + next_epoch = resp['result']['epoch'] + + @staticmethod + async def wait_for_next_epoch_if_soon(async_client: AsyncClient): + resp = await async_client.get_epoch_info(commitment=Confirmed) + if resp['result']['slotsInEpoch'] - resp['result']['slotIndex'] < 10: + await Waiter.wait_for_next_epoch(async_client) + + +@pytest.fixture +def waiter() -> Waiter: + return Waiter() diff --git a/stake-pool/py/tests/test_a_time_sensitive.py b/stake-pool/py/tests/test_a_time_sensitive.py index e32679832d8..78f707c7d0d 100644 --- a/stake-pool/py/tests/test_a_time_sensitive.py +++ b/stake-pool/py/tests/test_a_time_sensitive.py @@ -10,7 +10,7 @@ @pytest.mark.asyncio -async def test_increase_decrease_this_is_very_slow(async_client, validators, payer, stake_pool_addresses): +async def test_increase_decrease_this_is_very_slow(async_client, validators, payer, stake_pool_addresses, waiter): (stake_pool_address, validator_list_address) = stake_pool_addresses resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN) stake_rent_exemption = resp['result'] @@ -38,8 +38,8 @@ async def test_increase_decrease_this_is_very_slow(async_client, validators, pay assert validator.transient_stake_lamports == increase_amount + stake_rent_exemption assert validator.active_stake_lamports == 0 - print("Waiting for epoch to roll over, roughly 24 seconds") - await asyncio.sleep(24.0) + print("Waiting for epoch to roll over") + await waiter.wait_for_next_epoch(async_client) await update_stake_pool(async_client, payer, stake_pool_address) resp = await async_client.get_account_info(validator_list_address, commitment=Confirmed) @@ -64,8 +64,8 @@ async def test_increase_decrease_this_is_very_slow(async_client, validators, pay assert validator.transient_stake_lamports == decrease_amount assert validator.active_stake_lamports == increase_amount - decrease_amount - print("Waiting for epoch to roll over, roughly 24 seconds") - await asyncio.sleep(24.0) + print("Waiting for epoch to roll over") + await waiter.wait_for_next_epoch(async_client) await update_stake_pool(async_client, payer, stake_pool_address) resp = await async_client.get_account_info(validator_list_address, commitment=Confirmed) diff --git a/stake-pool/py/tests/test_bot_rebalance.py b/stake-pool/py/tests/test_bot_rebalance.py new file mode 100644 index 00000000000..1b28251b001 --- /dev/null +++ b/stake-pool/py/tests/test_bot_rebalance.py @@ -0,0 +1,79 @@ +"""Time sensitive test, so run it first out of the bunch.""" +import pytest +from solana.rpc.commitment import Confirmed +from spl.token.instructions import get_associated_token_address + +from stake.constants import STAKE_LEN +from stake_pool.state import StakePool, ValidatorList +from stake_pool.actions import deposit_sol + +from bot.rebalance import rebalance + + +ENDPOINT: str = "http://127.0.0.1:8899" + + +@pytest.mark.asyncio +async def test_rebalance_this_is_very_slow(async_client, validators, payer, stake_pool_addresses, waiter): + (stake_pool_address, validator_list_address) = stake_pool_addresses + resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN) + stake_rent_exemption = resp['result'] + increase_amount = 100_000_000 + deposit_amount = (increase_amount + stake_rent_exemption) * len(validators) + + resp = await async_client.get_account_info(stake_pool_address, commitment=Confirmed) + data = resp['result']['value']['data'] + stake_pool = StakePool.decode(data[0], data[1]) + token_account = get_associated_token_address(payer.public_key, stake_pool.pool_mint) + await deposit_sol(async_client, payer, stake_pool_address, token_account, deposit_amount) + + # Test case 1: Increase + await rebalance(ENDPOINT, stake_pool_address, payer, 0.0) + + # should only have minimum left + resp = await async_client.get_account_info(stake_pool.reserve_stake, commitment=Confirmed) + assert resp['result']['value']['lamports'] == stake_rent_exemption + 1 + + # should all be the same + resp = await async_client.get_account_info(validator_list_address, commitment=Confirmed) + data = resp['result']['value']['data'] + validator_list = ValidatorList.decode(data[0], data[1]) + for validator in validator_list.validators: + assert validator.active_stake_lamports == 0 + assert validator.transient_stake_lamports == increase_amount + stake_rent_exemption + + # Test case 2: Decrease + print('Waiting for next epoch') + await waiter.wait_for_next_epoch(async_client) + await rebalance(ENDPOINT, stake_pool_address, payer, deposit_amount / 1_000_000_000) + + # should still only have minimum left + rent exemptions from increase + resp = await async_client.get_account_info(stake_pool.reserve_stake, commitment=Confirmed) + reserve_lamports = resp['result']['value']['lamports'] + assert reserve_lamports == stake_rent_exemption * (1 + len(validator_list.validators)) + 1 + + # should all be decreasing now + resp = await async_client.get_account_info(validator_list_address, commitment=Confirmed) + data = resp['result']['value']['data'] + validator_list = ValidatorList.decode(data[0], data[1]) + for validator in validator_list.validators: + assert validator.active_stake_lamports == 0 + assert validator.transient_stake_lamports == increase_amount + + # Test case 3: Do nothing + print('Waiting for next epoch') + await waiter.wait_for_next_epoch(async_client) + await rebalance(ENDPOINT, stake_pool_address, payer, deposit_amount / 1_000_000_000) + + # should still only have minimum left + rent exemptions from increase + resp = await async_client.get_account_info(stake_pool.reserve_stake, commitment=Confirmed) + reserve_lamports = resp['result']['value']['lamports'] + assert reserve_lamports == stake_rent_exemption + deposit_amount + 1 + + # should all be decreasing now + resp = await async_client.get_account_info(validator_list_address, commitment=Confirmed) + data = resp['result']['value']['data'] + validator_list = ValidatorList.decode(data[0], data[1]) + for validator in validator_list.validators: + assert validator.active_stake_lamports == 0 + assert validator.transient_stake_lamports == 0 diff --git a/stake-pool/py/tests/test_deposit_withdraw_stake.py b/stake-pool/py/tests/test_deposit_withdraw_stake.py index 3341d0d9bfa..12340a44eac 100644 --- a/stake-pool/py/tests/test_deposit_withdraw_stake.py +++ b/stake-pool/py/tests/test_deposit_withdraw_stake.py @@ -1,16 +1,30 @@ +import asyncio import pytest +from typing import Tuple +from solana.rpc.async_api import AsyncClient from solana.rpc.commitment import Confirmed from solana.keypair import Keypair +from solana.publickey import PublicKey from spl.token.instructions import get_associated_token_address from stake.actions import create_stake, delegate_stake from stake.constants import STAKE_LEN -from stake_pool.actions import deposit_stake, withdraw_stake +from stake_pool.actions import deposit_stake, withdraw_stake, update_stake_pool from stake_pool.state import StakePool +async def prepare_stake( + async_client: AsyncClient, payer: Keypair, stake_pool_address: PublicKey, + validator: PublicKey, token_account: PublicKey, stake_amount: int +) -> Tuple[PublicKey, PublicKey]: + stake = Keypair() + await create_stake(async_client, payer, stake, payer.public_key, stake_amount) + await delegate_stake(async_client, payer, payer, stake.public_key, validator) + return (stake.public_key, validator) + + @pytest.mark.asyncio -async def test_deposit_withdraw_stake(async_client, validators, payer, stake_pool_addresses): +async def test_deposit_withdraw_stake(async_client, validators, payer, stake_pool_addresses, waiter): (stake_pool_address, validator_list_address) = stake_pool_addresses resp = await async_client.get_account_info(stake_pool_address, commitment=Confirmed) data = resp['result']['value']['data'] @@ -19,24 +33,35 @@ async def test_deposit_withdraw_stake(async_client, validators, payer, stake_poo resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN) stake_rent_exemption = resp['result'] + await waiter.wait_for_next_epoch_if_soon(async_client) + await update_stake_pool(async_client, payer, stake_pool_address) stake_amount = 1_000_000 - for validator in validators: - stake = Keypair() - await create_stake(async_client, payer, stake, payer.public_key, stake_amount) - await delegate_stake(async_client, payer, payer, stake.public_key, validator) - await deposit_stake(async_client, payer, stake_pool_address, validator, stake.public_key, token_account) + futures = [ + prepare_stake(async_client, payer, stake_pool_address, validator, token_account, stake_amount) + for validator in validators + ] + stakes = await asyncio.gather(*futures) + await waiter.wait_for_next_epoch(async_client) + await update_stake_pool(async_client, payer, stake_pool_address) + futures = [ + deposit_stake(async_client, payer, stake_pool_address, validator, stake, token_account) + for (stake, validator) in stakes + ] + stakes = await asyncio.gather(*futures) pool_token_balance = await async_client.get_token_account_balance(token_account, Confirmed) pool_token_balance = pool_token_balance['result']['value']['amount'] assert pool_token_balance == str((stake_amount + stake_rent_exemption) * len(validators)) + futures = [] for validator in validators: destination_stake = Keypair() - await withdraw_stake( + futures.append(withdraw_stake( async_client, payer, payer, destination_stake, stake_pool_address, validator, payer.public_key, token_account, stake_amount - ) + )) + await asyncio.gather(*futures) pool_token_balance = await async_client.get_token_account_balance(token_account, Confirmed) pool_token_balance = pool_token_balance['result']['value']['amount']