From 22062d9aa065a303531d786d9f3bde23a9045e60 Mon Sep 17 00:00:00 2001 From: SCM Date: Thu, 14 Nov 2024 12:13:32 +0000 Subject: [PATCH] More tests / bug fixes --- plutus_bench/mock.py | 22 +++++ plutus_bench/tx_tools.py | 2 +- tests/contracts/unrealistic_staking.py | 8 +- tests/stake.py | 14 ++- tests/test_stake.py | 121 +++++++++++++++++++++++-- 5 files changed, 153 insertions(+), 14 deletions(-) diff --git a/plutus_bench/mock.py b/plutus_bench/mock.py index fc7c0b3..802a3ec 100644 --- a/plutus_bench/mock.py +++ b/plutus_bench/mock.py @@ -211,6 +211,17 @@ def submit_tx(self, tx: Transaction): self.submit_tx_mock(tx) def submit_tx_mock(self, tx: Transaction): + def is_witnessed(address: Union[bytes, pycardano.Address], witness_set: pycardano.TransactionWitnessSet) -> bool: + if isinstance(address, bytes): + address = pycardano.Address.from_primitive(address) + staking_part = address.staking_part + if isinstance(staking_part, pycardano.ScriptHash): + scripts = (witness_set.plutus_v1_script or []) + (witness_set.plutus_v2_script or []) + (witness_set.plutus_v3_script or []) + return staking_part in [pycardano.plutus_script_hash(s) for s in scripts] + else: + raise NotImplementedError() + + for input in tx.transaction_body.inputs: utxo = self.get_utxo_from_txid(input.transaction_id, input.index) self.remove_utxo(utxo) @@ -251,6 +262,17 @@ def submit_tx_mock(self, tx: Transaction): self._pool_delegators[str(pool_id)].append( certificate.stake_credential.credential ) + for address in tx.transaction_body.withdraws or {}: + value = tx.transaction_body.withdraws[address] + stake_address = pycardano.Address.from_primitive(address) + assert is_witnessed(stake_address, tx.transaction_witness_set), f'Withdrawal from address {stake_address} is not witnessed' + assert str(stake_address) in self._reward_account, 'Address {stake_address} not registered' + rewards = self._reward_account[str(stake_address)]['delegation']['rewards'] + assert rewards == value, 'All rewards must be withdrawn. Requested {value} but account contains {rewards}' + self._reward_account[str(stake_address)]['delegation']['rewards'] == 0 + + + def submit_tx_cbor(self, cbor: Union[bytes, str]): return self.submit_tx(Transaction.from_cbor(cbor)) diff --git a/plutus_bench/tx_tools.py b/plutus_bench/tx_tools.py index 06c0ef4..61b552f 100644 --- a/plutus_bench/tx_tools.py +++ b/plutus_bench/tx_tools.py @@ -404,7 +404,7 @@ def generate_script_contexts_resolved( certificate_redeemer = as_redeemer( next( r - for r in tx.transaction_witness_set.redeemer + for r in tx.transaction_witness_set.redeemer or [] if r.index == i and r.tag == RedeemerTag.CERTIFICATE ), tx.transaction_witness_set.redeemer, diff --git a/tests/contracts/unrealistic_staking.py b/tests/contracts/unrealistic_staking.py index 81bc7d1..195355d 100644 --- a/tests/contracts/unrealistic_staking.py +++ b/tests/contracts/unrealistic_staking.py @@ -1,9 +1,7 @@ from opshin.prelude import * -def validator( - address: Address, datum: BuiltinData, redeemer: BuiltinData, context: ScriptContext -) -> None: +def validator(address: Address, redeemer: BuiltinData, context: ScriptContext) -> None: purpose = context.purpose if isinstance(purpose, Certifying): return None # Do whatever you like with certifiying @@ -12,6 +10,8 @@ def validator( paid_to_address = all_tokens_locked_at_address( context.tx_info.outputs, address, Token(b"", b"") ) - assert paid_to_address >= 2 * withdrawal_amount + assert ( + paid_to_address >= 2 * withdrawal_amount + ), "Insufficient rewards to address" else: assert False, "not a valid purpose" diff --git a/tests/stake.py b/tests/stake.py index 7ce14ae..af37b9a 100644 --- a/tests/stake.py +++ b/tests/stake.py @@ -13,6 +13,8 @@ def register_and_delegate( plutus_script: pycardano.PlutusV2Script, pool_id: PoolId, context: ChainContext, + reverse_cert_order=False, + add_certificate_script=True, ): delegator_vkey_hash = delegator_skey.to_verification_key().hash() @@ -36,9 +38,13 @@ def register_and_delegate( builder = pycardano.TransactionBuilder(context) builder.add_input_address(delegator_address) - builder.certificates = [stake_registration, stake_delegation] - redeemer = pycardano.Redeemer(0) - builder.add_certificate_script(plutus_script, redeemer=redeemer) + if reverse_cert_order: + builder.certificates = [stake_delegation, stake_registration] + else: + builder.certificates = [stake_registration, stake_delegation] + if add_certificate_script: + redeemer = pycardano.Redeemer(0) + builder.add_certificate_script(plutus_script, redeemer=redeemer) tx = builder.build_and_sign( signing_keys=[delegator_skey], change_address=script_payment_address, @@ -77,3 +83,5 @@ def withdraw( builder.add_withdrawal_script(plutus_script, redeemer=redeemer) builder.add_output(pycardano.TransactionOutput(recipient_address, recipient_amount)) tx = builder.build_and_sign([delegator_skey], script_payment_address) + context.submit_tx(tx) + return dict(stake_address=stake_address) diff --git a/tests/test_stake.py b/tests/test_stake.py index 30c2c05..d72d937 100644 --- a/tests/test_stake.py +++ b/tests/test_stake.py @@ -10,10 +10,23 @@ from tests.stake import register_and_delegate, withdraw from pycardano.crypto.bech32 import decode from opshin import build +from opshin.ledger.api_v2 import ( + Address as Address, + PubKeyCredential, + PubKeyHash, + NoStakingCredential, +) own_path = pathlib.Path(__file__) +def as_ledger_address(address: pycardano.Address) -> Address: + return Address( + PubKeyCredential(PubKeyHash(address.payment_part.payload)), + NoStakingCredential(), + ) + + def test_register_and_delegate(): api = MockFrostApi() context = MockChainContext(api=api) @@ -22,12 +35,52 @@ def test_register_and_delegate(): staking_user.fund(100_000_000_000) script_path = own_path.parent / "contracts/unrealistic_staking.py" - plutus_script = build(script_path, bytes(staking_user.verification_key.hash())) + plutus_script = build(script_path, as_ledger_address(staking_user.address)) register_and_delegate( staking_user.signing_key, plutus_script, stake_pool.pool_id, context ) +def test_register_and_delegate_wrong_order(): + api = MockFrostApi() + context = MockChainContext(api=api) + staking_user = MockUser(api) + stake_pool = MockPool(api) + + staking_user.fund(100_000_000_000) + script_path = own_path.parent / "contracts/unrealistic_staking.py" + plutus_script = build(script_path, as_ledger_address(staking_user.address)) + pytest.raises( + TransactionFailedException, + register_and_delegate, + staking_user.signing_key, + plutus_script, + stake_pool.pool_id, + context, + reverse_cert_order=True, + ) + + +def test_register_and_delegate_no_script(): + api = MockFrostApi() + context = MockChainContext(api=api) + staking_user = MockUser(api) + stake_pool = MockPool(api) + + staking_user.fund(100_000_000_000) + script_path = own_path.parent / "contracts/unrealistic_staking.py" + plutus_script = build(script_path, as_ledger_address(staking_user.address)) + pytest.raises( + ValueError, + register_and_delegate, + staking_user.signing_key, + plutus_script, + stake_pool.pool_id, + context, + add_certificate_script=False, + ) + + def test_withdraw(): api = MockFrostApi() context = MockChainContext(api=api) @@ -38,7 +91,7 @@ def test_withdraw(): staking_user.fund(100_000_000_000) script_path = own_path.parent / "contracts/unrealistic_staking.py" - plutus_script = build(script_path, bytes(recipient_user.verification_key.hash())) + plutus_script = build(script_path, as_ledger_address(recipient_user.address)) stake_info = register_and_delegate( staking_user.signing_key, plutus_script, stake_pool.pool_id, context ) @@ -52,10 +105,66 @@ def test_withdraw(): context, ) - stake_address = stake_info["stake_address"] - script_payment_address = stake_info["script_payment_address"] + +def test_withdraw_rewards(): + api = MockFrostApi() + context = MockChainContext(api=api) + staking_user = MockUser(api) + recipient_user = MockUser(api) + stake_pool = MockPool(api) + + staking_user.fund(100_000_000_000) + script_path = own_path.parent / "contracts/unrealistic_staking.py" + plutus_script = build(script_path, as_ledger_address(recipient_user.address)) + stake_info = register_and_delegate( + staking_user.signing_key, plutus_script, stake_pool.pool_id, context + ) + + api.distribute_rewards(10_000_000_000) + + # Withdraw + withdraw( + recipient_user.address, + 60_000_000_000, + staking_user.signing_key, + plutus_script, + context, + ) + + +def test_withdraw_script_failure(): + api = MockFrostApi() + context = MockChainContext(api=api) + staking_user = MockUser(api) + recipient_user = MockUser(api) + stake_pool = MockPool(api) + + staking_user.fund(100_000_000_000) + + script_path = own_path.parent / "contracts/unrealistic_staking.py" + plutus_script = build(script_path, as_ledger_address(recipient_user.address)) + stake_info = register_and_delegate( + staking_user.signing_key, plutus_script, stake_pool.pool_id, context + ) + + api.distribute_rewards(10_000_000_000) + + # Fails if recipient recieves less than double the reward amount + pytest.raises( + TransactionFailedException, + withdraw, + recipient_user.address, + 19_000_000_000, + staking_user.signing_key, + plutus_script, + context, + ) if __name__ == "__main__": - test_register_and_delegate() - test_withdraw() + # test_register_and_delegate() + # test_register_and_delegate_wrong_order() + # test_register_and_delegate_no_script() + # test_withdraw() + test_withdraw_rewards() + # test_withdraw_script_failure()