stake-pool-py: Add more waits to solidify tests (#2882)

Author: joncinque

stake-pool/py/stake/state.py

@@ -2,9 +2,10 @@
 
 from enum import IntEnum
 from typing import NamedTuple, Dict
-from construct import Container, Struct, Int64ul  # type: ignore
+from construct import Container, Struct, Float64l, Int32ul, Int64ul  # type: ignore
 
 from solana.publickey import PublicKey
+from solana.utils.helpers import decode_byte_string
 from solana._layouts.shared import PUBLIC_KEY_LAYOUT
 
 
@@ -46,6 +47,29 @@ class StakeAuthorize(IntEnum):
     WITHDRAWER = 1
 
 
+class StakeStateType(IntEnum):
+    """Stake State Types."""
+    UNINITIALIZED = 0
+    INITIALIZED = 1
+    STAKE = 2
+    REWARDS_POOL = 3
+
+
+class StakeState(NamedTuple):
+    state_type: StakeStateType
+    state: Container
+
+    """Stake state."""
+    @classmethod
+    def decode(cls, data: str, encoding: str):
+        data_bytes = decode_byte_string(data, encoding)
+        parsed = STAKE_STATE_LAYOUT.parse(data_bytes)
+        return StakeState(
+            state_type=parsed['state_type'],
+            state=parsed['state'],
+        )
+
+
 LOCKUP_LAYOUT = Struct(
     "unix_timestamp" / Int64ul,
     "epoch" / Int64ul,
@@ -57,3 +81,51 @@ class StakeAuthorize(IntEnum):
     "staker" / PUBLIC_KEY_LAYOUT,
     "withdrawer" / PUBLIC_KEY_LAYOUT,
 )
+
+META_LAYOUT = Struct(
+    "rent_exempt_reserve" / Int64ul,
+    "authorized" / AUTHORIZED_LAYOUT,
+    "lockup" / LOCKUP_LAYOUT,
+)
+
+META_LAYOUT = Struct(
+    "rent_exempt_reserve" / Int64ul,
+    "authorized" / AUTHORIZED_LAYOUT,
+    "lockup" / LOCKUP_LAYOUT,
+)
+
+DELEGATION_LAYOUT = Struct(
+    "voter_pubkey" / PUBLIC_KEY_LAYOUT,
+    "stake" / Int64ul,
+    "activation_epoch" / Int64ul,
+    "deactivation_epoch" / Int64ul,
+    "warmup_cooldown_rate" / Float64l,
+)
+
+STAKE_LAYOUT = Struct(
+    "delegation" / DELEGATION_LAYOUT,
+    "credits_observed" / Int64ul,
+)
+
+STAKE_AND_META_LAYOUT = Struct(
+    "meta" / META_LAYOUT,
+    "stake" / STAKE_LAYOUT,
+)
+
+STAKE_STATE_LAYOUT = Struct(
+    "state_type" / Int32ul,
+    "state" / STAKE_AND_META_LAYOUT,
+    # NOTE: This can be done better, but was mainly needed for testing. Ideally,
+    # we would have something like:
+    #
+    # Switch(
+    #     lambda this: this.state,
+    #     {
+    #         StakeStateType.UNINITIALIZED: Pass,
+    #         StakeStateType.INITIALIZED: META_LAYOUT,
+    #         StakeStateType.STAKE: STAKE_AND_META_LAYOUT,
+    #     }
+    # ),
+    #
+    # Unfortunately, it didn't work.
+)

stake-pool/py/tests/conftest.py

@@ -42,28 +42,28 @@ def solana_test_validator():
 def validators(event_loop, async_client, payer) -> List[PublicKey]:
     num_validators = 3
     validators = []
-    futures = []
     for i in range(num_validators):
         vote = Keypair()
         node = Keypair()
-        futures.append(create_vote(async_client, payer, vote, node, payer.public_key, payer.public_key, 10))
+        event_loop.run_until_complete(
+            create_vote(async_client, payer, vote, node, payer.public_key, payer.public_key, 10)
+        )
         validators.append(vote.public_key)
-    event_loop.run_until_complete(asyncio.gather(*futures))
     return validators
 
 
 @pytest.fixture
-def stake_pool_addresses(event_loop, async_client, payer, validators) -> Tuple[PublicKey, PublicKey]:
+def stake_pool_addresses(event_loop, async_client, payer, validators, waiter) -> Tuple[PublicKey, PublicKey]:
     fee = Fee(numerator=1, denominator=1000)
     referral_fee = 20
+    event_loop.run_until_complete(waiter.wait_for_next_epoch_if_soon(async_client))
     stake_pool_addresses = event_loop.run_until_complete(
         create_all(async_client, payer, fee, referral_fee)
     )
-    futures = [
-        add_validator_to_pool(async_client, payer, stake_pool_addresses[0], validator)
-        for validator in validators
-    ]
-    event_loop.run_until_complete(asyncio.gather(*futures))
+    for validator in validators:
+        event_loop.run_until_complete(
+            add_validator_to_pool(async_client, payer, stake_pool_addresses[0], validator)
+        )
     return stake_pool_addresses
 
 

stake-pool/py/tests/test_deposit_withdraw_stake.py

@@ -5,6 +5,7 @@
 
 from stake.actions import create_stake, delegate_stake
 from stake.constants import STAKE_LEN
+from stake.state import StakeState
 from stake_pool.actions import deposit_stake, withdraw_stake, update_stake_pool
 from stake_pool.state import StakePool
 
@@ -15,26 +16,26 @@ async def test_deposit_withdraw_stake(async_client, validators, payer, stake_poo
     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)
-
-    resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN)
-    stake_rent_exemption = resp['result']
-    waited = await waiter.wait_for_next_epoch_if_soon(async_client)
-    if waited:
-        await update_stake_pool(async_client, payer, stake_pool_address)
-
     validator = next(iter(validators))
     stake_amount = 1_000_000
     stake = Keypair()
     await create_stake(async_client, payer, stake, payer.public_key, stake_amount)
     stake = stake.public_key
     await delegate_stake(async_client, payer, payer, stake, validator)
+    resp = await async_client.get_account_info(stake, commitment=Confirmed)
+    data = resp['result']['value']['data']
+    stake_state = StakeState.decode(data[0], data[1])
+    print(stake_state)
+
     await waiter.wait_for_next_epoch(async_client)
+
     await update_stake_pool(async_client, payer, stake_pool_address)
+    token_account = get_associated_token_address(payer.public_key, stake_pool.pool_mint)
     await deposit_stake(async_client, payer, stake_pool_address, validator, stake, token_account)
-
     pool_token_balance = await async_client.get_token_account_balance(token_account, Confirmed)
     pool_token_balance = pool_token_balance['result']['value']['amount']
+    resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN)
+    stake_rent_exemption = resp['result']
     assert pool_token_balance == str(stake_amount + stake_rent_exemption)
 
     destination_stake = Keypair()