From f2bd0edc5151e2adb8a4a5d2b427eb5187500429 Mon Sep 17 00:00:00 2001
From: Alex Grintsvayg <git@grin.io>
Date: Mon, 17 Jun 2019 18:29:17 -0400
Subject: [PATCH] add option to only use confirmed utxos

---
 tests/client_tests/unit/test_coinselection.py | 10 ++++++++
 tests/client_tests/unit/test_transaction.py   |  4 ++--
 torba/client/baseledger.py                    |  4 +++-
 torba/client/coinselection.py                 | 23 ++++++++++++++++++-
 4 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/tests/client_tests/unit/test_coinselection.py b/tests/client_tests/unit/test_coinselection.py
index b74dcc768..8f3d3568a 100644
--- a/tests/client_tests/unit/test_coinselection.py
+++ b/tests/client_tests/unit/test_coinselection.py
@@ -85,6 +85,16 @@ class TestCoinSelectionTests(BaseSelectionTestCase):
         match = selector.select()
         self.assertEqual([5*CENT], [c.txo.amount for c in match])
 
+    def test_prefer_confirmed_strategy(self):
+        utxo_pool = self.estimates(
+            utxo(11*CENT, height=5),
+            utxo(11*CENT, height=0),
+            utxo(11*CENT, height=-2),
+            utxo(11*CENT, height=5),
+        )
+        selector = CoinSelector(utxo_pool, 20*CENT, 0)
+        match = selector.select("confirmed_only")
+        self.assertEqual([5,5], [c.txo.tx_ref.height for c in match])
 
 class TestOfficialBitcoinCoinSelectionTests(BaseSelectionTestCase):
 
diff --git a/tests/client_tests/unit/test_transaction.py b/tests/client_tests/unit/test_transaction.py
index 800c05322..1839aea06 100644
--- a/tests/client_tests/unit/test_transaction.py
+++ b/tests/client_tests/unit/test_transaction.py
@@ -14,8 +14,8 @@ FEE_PER_BYTE = 50
 FEE_PER_CHAR = 200000
 
 
-def get_output(amount=CENT, pubkey_hash=NULL_HASH):
-    return ledger_class.transaction_class() \
+def get_output(amount=CENT, pubkey_hash=NULL_HASH, height=-2):
+    return ledger_class.transaction_class(height=height) \
         .add_outputs([ledger_class.transaction_class.output_class.pay_pubkey_hash(amount, pubkey_hash)]) \
         .outputs[0]
 
diff --git a/torba/client/baseledger.py b/torba/client/baseledger.py
index 8775d6f24..c6a5f8a68 100644
--- a/torba/client/baseledger.py
+++ b/torba/client/baseledger.py
@@ -140,6 +140,8 @@ class BaseLedger(metaclass=LedgerRegistry):
         self._header_processing_lock = asyncio.Lock()
         self._address_update_locks: Dict[str, asyncio.Lock] = {}
 
+        self.coin_selection_strategy = None
+
     @classmethod
     def get_id(cls):
         return '{}_{}'.format(cls.symbol.lower(), cls.network_name.lower())
@@ -212,7 +214,7 @@ class BaseLedger(metaclass=LedgerRegistry):
                 txos, amount,
                 self.transaction_class.output_class.pay_pubkey_hash(COIN, NULL_HASH32).get_fee(self)
             )
-            spendables = selector.select()
+            spendables = selector.select(self.coin_selection_strategy)
             if spendables:
                 await self.reserve_outputs(s.txo for s in spendables)
             return spendables
diff --git a/torba/client/coinselection.py b/torba/client/coinselection.py
index 119a6e11f..e0d0238b6 100644
--- a/torba/client/coinselection.py
+++ b/torba/client/coinselection.py
@@ -5,6 +5,12 @@ from torba.client import basetransaction
 
 MAXIMUM_TRIES = 100000
 
+STRATEGIES = []
+
+def strategy(method):
+    STRATEGIES.append(method.__name__)
+    return method
+
 
 class CoinSelector:
 
@@ -20,17 +26,30 @@ class CoinSelector:
         if seed is not None:
             self.random.seed(seed, version=1)
 
-    def select(self) -> List[basetransaction.BaseOutputEffectiveAmountEstimator]:
+    def select(self, strategy: str = None) -> List[basetransaction.BaseOutputEffectiveAmountEstimator]:
         if not self.txos:
             return []
         if self.target > self.available:
             return []
+        if strategy is not None:
+            return getattr(self, strategy)()
         return (
             self.branch_and_bound() or
             self.closest_match() or
             self.random_draw()
         )
 
+    @strategy
+    def confirmed_only(self) -> List[basetransaction.BaseOutputEffectiveAmountEstimator]:
+        self.txos = [t for t in self.txos if t.txo.tx_ref.height > 0] or self.txos
+        self.available = sum(c.effective_amount for c in self.txos)
+        return (
+                self.branch_and_bound() or
+                self.closest_match() or
+                self.random_draw()
+        )
+
+    @strategy
     def branch_and_bound(self) -> List[basetransaction.BaseOutputEffectiveAmountEstimator]:
         # see bitcoin implementation for more info:
         # https://github.com/bitcoin/bitcoin/blob/master/src/wallet/coinselection.cpp
@@ -89,6 +108,7 @@ class CoinSelector:
 
         return []
 
+    @strategy
     def closest_match(self) -> List[basetransaction.BaseOutputEffectiveAmountEstimator]:
         """ Pick one UTXOs that is larger than the target but with the smallest change. """
         target = self.target + self.cost_of_change
@@ -101,6 +121,7 @@ class CoinSelector:
                     smallest_change, best_match = change, txo
         return [best_match] if best_match else []
 
+    @strategy
     def random_draw(self) -> List[basetransaction.BaseOutputEffectiveAmountEstimator]:
         """ Accumulate UTXOs at random until there is enough to cover the target. """
         target = self.target + self.cost_of_change