Merge pull request #16 from Merkleize/fix-wait_for_spend

Fix behavior of wait_for_spend so it tracks known StandardP2TR outputs
Merkleize · Apr 29, 2024 · 9d8861d · 9d8861d
2 parents 677035f + cea56fd
commit 9d8861d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 18 deletions.
diff --git a/matt/manager.py b/matt/manager.py
@@ -599,28 +599,29 @@ def wait_for_spend(self, instances: Union[ContractInstance, List[ContractInstanc
                     out_contract = clause_output.next_contract
                     new_instance = ContractInstance(out_contract)
 
-                    if isinstance(out_contract, (P2TR, OpaqueP2TR, StandardP2TR)):
-                        continue  # nothing to do, will not track this output
-                    elif isinstance(out_contract, StandardAugmentedP2TR):
-                        if clause_output.next_state is None:
-                            raise ValueError("Missing data for augmented output")
-                        new_instance.data = clause_output.next_state.encode()
-                        new_instance.data_expanded = clause_output.next_state
-                    else:
-                        raise ValueError("Unsupported contract type")
+                    if isinstance(out_contract, (StandardP2TR, StandardAugmentedP2TR)):
+                        if isinstance(out_contract, StandardAugmentedP2TR):
+                            if clause_output.next_state is None:
+                                raise ValueError("Missing data for augmented output")
+                            new_instance.data = clause_output.next_state.encode()
+                            new_instance.data_expanded = clause_output.next_state
 
-                    new_instance.last_height = instance.last_height
+                        new_instance.last_height = instance.last_height
 
-                    new_instance.outpoint = COutPoint(int(tx.hash, 16), output_index)
-                    new_instance.funding_tx = tx
-                    new_instance.status = ContractInstanceStatus.FUNDED
+                        new_instance.outpoint = COutPoint(int(tx.hash, 16), output_index)
+                        new_instance.funding_tx = tx
+                        new_instance.status = ContractInstanceStatus.FUNDED
 
-                    out_contracts[output_index] = new_instance
+                        out_contracts[output_index] = new_instance
 
-                    next_instances.append(new_instance)
+                        next_instances.append(new_instance)
+                    elif isinstance(out_contract, (P2TR, OpaqueP2TR)):
+                        continue  # nothing to do, will not track this output
+                    else:
+                        raise ValueError("Unsupported contract type")
                 instance.next = next_instances
 
-        result = list(out_contracts.values())
+        result = [out_contracts[idx] for idx in sorted(out_contracts.keys())]
         for instance in result:
             self.add_instance(instance)
         return result

diff --git a/tests/test_vault.py b/tests/test_vault.py
@@ -1,7 +1,7 @@
 from typing import Tuple
 import pytest
 
-from examples.vault.vault_contracts import Vault
+from examples.vault.vault_contracts import Vault, Unvaulting
 
 from matt.btctools import key
 from matt.btctools.auth_proxy import AuthServiceProxy, JSONRPCException
@@ -183,7 +183,12 @@ def test_vault_trigger_with_revault_and_withdraw(vault_specs: VaultSpecs, rpc: A
             {**args, "sig": sigs[i]}
         ))
 
-    [U_inst] = manager.spend_and_wait([V_inst_1, V_inst_2, V_inst_3], spend_tx)
+    [U_inst, V_revault_inst] = manager.spend_and_wait([V_inst_1, V_inst_2, V_inst_3], spend_tx)
+
+    assert isinstance(U_inst.contract, Unvaulting)
+    assert isinstance(V_revault_inst.contract, Vault)
+    assert manager.instances.index(U_inst) >= 0
+    assert manager.instances.index(V_revault_inst) >= 0
 
     report.write(vault_description, format_tx_markdown(spend_tx, "Trigger (with revault) [3 vault inputs]"))