More tests / bug fixes

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))

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,

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"

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)

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()