Merge #406: Unit tests and modularisation of YieldGeneratorBasic

6a0e742836 Add basic unit tests for YieldGeneratorBasic.
8b1e24e20a Make yg algorithms easier to define.

jmclient/jmclient/yieldgenerator.py

@@ -17,7 +17,6 @@ from .wallet_utils import open_test_wallet_maybe, get_wallet_path
 
 jlog = get_log()
 
-MAX_MIX_DEPTH = 5
 
 class YieldGenerator(Maker):
     """A maker for the purposes of generating a yield from held
@@ -77,12 +76,11 @@ class YieldGeneratorBasic(YieldGenerator):
         super(YieldGeneratorBasic,self).__init__(wallet)
 
     def create_my_orders(self):
-        mix_balance = self.wallet.get_balance_by_mixdepth(verbose=False)
+        mix_balance = self.get_available_mixdepths()
         if len([b for m, b in iteritems(mix_balance) if b > 0]) == 0:
             jlog.error('do not have any coins left')
             return []
 
-        # print mix_balance
         max_mix = max(mix_balance, key=mix_balance.get)
         f = '0'
         if self.ordertype in ('reloffer', 'swreloffer'):
@@ -113,23 +111,24 @@ class YieldGeneratorBasic(YieldGenerator):
 
     def oid_to_order(self, offer, amount):
         total_amount = amount + offer["txfee"]
-        mix_balance = self.wallet.get_balance_by_mixdepth()
-        max_mix = max(mix_balance, key=mix_balance.get)
+        mix_balance = self.get_available_mixdepths()
 
-        filtered_mix_balance = [m
-                                for m in iteritems(mix_balance)
-                                if m[1] >= total_amount]
+        filtered_mix_balance = {m: b
+                                for m, b in iteritems(mix_balance)
+                                if b >= total_amount}
         if not filtered_mix_balance:
             return None, None, None
         jlog.debug('mix depths that have enough = ' + str(filtered_mix_balance))
-        filtered_mix_balance = sorted(filtered_mix_balance, key=lambda x: x[0])
-        mixdepth = filtered_mix_balance[0][0]
+        mixdepth = self.select_input_mixdepth(filtered_mix_balance, offer, amount)
+        if mixdepth is None:
+            return None, None, None
         jlog.info('filling offer, mixdepth=' + str(mixdepth) + ', amount=' + str(amount))
 
-        # mixdepth is the chosen depth we'll be spending from
-        cj_addr = self.wallet.get_internal_addr(
-            (mixdepth + 1) % (self.wallet.mixdepth + 1),
-            jm_single().bc_interface)
+        cj_addr = self.select_output_address(mixdepth, offer, amount)
+        if cj_addr is None:
+            return None, None, None
+        jlog.info('sending output to address=' + str(cj_addr))
+
         change_addr = self.wallet.get_internal_addr(mixdepth,
                                                     jm_single().bc_interface)
 
@@ -165,6 +164,28 @@ class YieldGeneratorBasic(YieldGenerator):
                 confirm_time / 60.0, 2), ''])
         return self.on_tx_unconfirmed(offer, txid, None)
 
+    def get_available_mixdepths(self):
+        """Returns the mixdepth/balance dict from the wallet that contains
+        all available inputs for offers."""
+        return self.wallet.get_balance_by_mixdepth(verbose=False)
+
+    def select_input_mixdepth(self, available, offer, amount):
+        """Returns the mixdepth from which the given order should spend the
+        inputs.  available is a mixdepth/balance dict of all the mixdepths
+        that can be chosen from, i.e. have enough balance.  If there is no
+        suitable input, the function can return None to abort the order."""
+        available = sorted(iteritems(available), key=lambda entry: entry[0])
+        return available[0][0]
+
+    def select_output_address(self, input_mixdepth, offer, amount):
+        """Returns the address to which the mixed output should be sent for
+        an order spending from the given input mixdepth.  Can return None if
+        there is no suitable output, in which case the order is
+        aborted."""
+        cjoutmix = (input_mixdepth + 1) % (self.wallet.mixdepth + 1)
+        return self.wallet.get_internal_addr(cjoutmix, jm_single().bc_interface)
+
+
 def ygmain(ygclass, txfee=1000, cjfee_a=200, cjfee_r=0.002, ordertype='swreloffer',
            nickserv_password='', minsize=100000, gaplimit=6):
     import sys

jmclient/test/test_yieldgenerator.py

@@ -0,0 +1,204 @@
+from __future__ import (absolute_import, division,
+                        print_function, unicode_literals)
+from builtins import * # noqa: F401
+
+import unittest
+
+from jmclient import load_program_config, jm_single,\
+    SegwitLegacyWallet, VolatileStorage, YieldGeneratorBasic, get_network
+
+
+class CustomUtxoWallet(SegwitLegacyWallet):
+    """A wallet instance that makes it easy to do the tasks we need for
+    the tests.  It has convenience methods to add UTXOs with pre-defined
+    balances for all mixdepths, and to verify if a given UTXO or address
+    corresponds to the mixdepth it should."""
+
+    def __init__(self, balances):
+        """Creates the wallet, setting the balances of the mixdepths
+        as given by the array.  (And the number of mixdepths from the array
+        elements."""
+
+        load_program_config()
+
+        storage = VolatileStorage()
+        super(CustomUtxoWallet, self).initialize(storage, get_network(),
+                                                 max_mixdepth=len(balances)-1)
+        super(CustomUtxoWallet, self).__init__(storage)
+
+        for m, b in enumerate(balances):
+            self.add_utxo_at_mixdepth(m, b)
+
+    def add_utxo_at_mixdepth(self, mixdepth, balance):
+        tx = {'outs': [{'script': self.get_internal_script(mixdepth),
+                        'value': balance}]}
+        # We need to generate a fake "txid" that has to be unique for all the
+        # UTXOs that are added to the wallet.  For that, we simply use the
+        # script, and make it fit the required length (32 bytes).
+        txid = tx['outs'][0]['script'] + b'x' * 32
+        txid = txid[:32]
+        self.add_new_utxos_(tx, txid)
+
+    def assert_utxos_from_mixdepth(self, utxos, expected):
+        """Asserts that the list of UTXOs (as returned from UTXO selection
+        of the wallet) is all from the given expected mixdepth."""
+
+        for u in utxos.values():
+            assert self.get_addr_mixdepth(u['address']) == expected
+
+
+def create_yg_basic(balances, txfee=0, cjfee_a=0, cjfee_r=0,
+                    ordertype='swabsoffer', minsize=0):
+    """Constructs a YieldGeneratorBasic instance with a fake wallet.  The
+    wallet will have the given balances at mixdepths, and the offer params
+    will be set as given here."""
+
+    wallet = CustomUtxoWallet(balances)
+    offerconfig = (txfee, cjfee_a, cjfee_r, ordertype, minsize)
+
+    yg = YieldGeneratorBasic(wallet, offerconfig)
+
+    # We don't need any of the logic from Maker, including the waiting
+    # loop.  Just stop it, so that it does not linger around and create
+    # unclean-reactor failures.
+    if yg.sync_wait_loop.running:
+        yg.sync_wait_loop.stop()
+
+    return yg
+
+
+class CreateMyOrdersTests(unittest.TestCase):
+    """Unit tests for YieldGeneratorBasic.create_my_orders."""
+
+    def test_no_coins(self):
+        yg = create_yg_basic([0] * 3)
+        self.assertEqual(yg.create_my_orders(), [])
+
+    def test_abs_fee(self):
+        jm_single().DUST_THRESHOLD = 10
+        yg = create_yg_basic([0, 2000000, 1000000], txfee=1000, cjfee_a=10,
+                             ordertype='swabsoffer', minsize=100000)
+        self.assertEqual(yg.create_my_orders(), [
+          {'oid': 0,
+           'ordertype': 'swabsoffer',
+           'minsize': 100000,
+           'maxsize': 1999000,
+           'txfee': 1000,
+           'cjfee': '1010'},
+        ])
+
+    def test_rel_fee(self):
+        jm_single().DUST_THRESHOLD = 10
+        yg = create_yg_basic([0, 2000000, 1000000], txfee=1000, cjfee_r=0.1,
+                             ordertype='swreloffer', minsize=10)
+        self.assertEqual(yg.create_my_orders(), [
+          {'oid': 0,
+           'ordertype': 'swreloffer',
+           'minsize': 15000,
+           'maxsize': 1999000,
+           'txfee': 1000,
+           'cjfee': 0.1},
+        ])
+
+    def test_dust_threshold(self):
+        jm_single().DUST_THRESHOLD = 1000
+        yg = create_yg_basic([0, 2000000, 1000000], txfee=10, cjfee_a=10,
+                             ordertype='swabsoffer', minsize=100000)
+        self.assertEqual(yg.create_my_orders(), [
+          {'oid': 0,
+           'ordertype': 'swabsoffer',
+           'minsize': 100000,
+           'maxsize': 1999000,
+           'txfee': 10,
+           'cjfee': '20'},
+        ])
+
+    def test_minsize_above_maxsize(self):
+        jm_single().DUST_THRESHOLD = 10
+        yg = create_yg_basic([0, 20000, 10000], txfee=1000, cjfee_a=10,
+                             ordertype='swabsoffer', minsize=100000)
+        self.assertEqual(yg.create_my_orders(), [])
+
+
+class OidToOrderTests(unittest.TestCase):
+    """Tests YieldGeneratorBasic.oid_to_order."""
+
+    def call_oid_to_order(self, yg, amount):
+        """Calls oid_to_order on the given yg instance.  It passes the
+        txfee and abs fee from yg as offer."""
+        offer = {'txfee': yg.txfee,
+                 'cjfee': str(yg.cjfee_a),
+                 'ordertype': 'swabsoffer'}
+        return yg.oid_to_order(offer, amount)
+
+    def test_not_enough_balance(self):
+        yg = create_yg_basic([100], txfee=0, cjfee_a=10)
+        self.assertEqual(self.call_oid_to_order(yg, 1000), (None, None, None))
+
+    def test_chooses_single_utxo(self):
+        jm_single().DUST_THRESHOLD = 10
+        yg = create_yg_basic([10, 1000, 2000])
+        utxos, cj_addr, change_addr = self.call_oid_to_order(yg, 500)
+        self.assertEqual(len(utxos), 1)
+        yg.wallet.assert_utxos_from_mixdepth(utxos, 1)
+        self.assertEqual(yg.wallet.get_addr_mixdepth(cj_addr), 2)
+        self.assertEqual(yg.wallet.get_addr_mixdepth(change_addr), 1)
+
+    def test_not_enough_balance_with_dust_threshold(self):
+        # 410 is exactly the size of the change output.  So it will be
+        # right at the dust threshold.  The wallet won't be able to find
+        # any extra inputs, though.
+        jm_single().DUST_THRESHOLD = 410
+        yg = create_yg_basic([10, 1000, 10], txfee=100, cjfee_a=10)
+        self.assertEqual(self.call_oid_to_order(yg, 500), (None, None, None))
+
+    def test_extra_with_dust_threshold(self):
+        # The output will be right at the dust threshold, so that we will
+        # need to include the extra_utxo from the wallet as well to get
+        # over the threshold.
+        jm_single().DUST_THRESHOLD = 410
+        yg = create_yg_basic([10, 1000, 10], txfee=100, cjfee_a=10)
+        yg.wallet.add_utxo_at_mixdepth(1, 500)
+        utxos, cj_addr, change_addr = self.call_oid_to_order(yg, 500)
+        self.assertEqual(len(utxos), 2)
+        yg.wallet.assert_utxos_from_mixdepth(utxos, 1)
+        self.assertEqual(yg.wallet.get_addr_mixdepth(cj_addr), 2)
+        self.assertEqual(yg.wallet.get_addr_mixdepth(change_addr), 1)
+
+
+class OfferReannouncementTests(unittest.TestCase):
+    """Tests offer reannouncement logic from on_tx_unconfirmed."""
+
+    def call_on_tx_unconfirmed(self, yg):
+        """Calls yg.on_tx_unconfirmed with fake arguments."""
+        return yg.on_tx_unconfirmed({'cjaddr': 'addr'}, 'txid', [])
+
+    def create_yg_and_offer(self, maxsize):
+        """Constructs a fake yg instance that has an offer with the given
+        maxsize.  Returns it together with the offer."""
+        jm_single().DUST_THRESHOLD = 10
+        yg = create_yg_basic([100 + maxsize], txfee=100, ordertype='swabsoffer')
+        offers = yg.create_my_orders()
+        self.assertEqual(len(offers), 1)
+        self.assertEqual(offers[0]['maxsize'], maxsize)
+        return yg, offers[0]
+
+    def test_no_new_offers(self):
+        yg = create_yg_basic([0] * 3)
+        yg.offerlist = [{'oid': 0}]
+        self.assertEqual(self.call_on_tx_unconfirmed(yg), ([0], []))
+
+    def test_no_old_offers(self):
+        yg, offer = self.create_yg_and_offer(100)
+        yg.offerlist = []
+        self.assertEqual(self.call_on_tx_unconfirmed(yg), ([], [offer]))
+
+    def test_offer_unchanged(self):
+        yg, offer = self.create_yg_and_offer(100)
+        yg.offerlist = [offer]
+        self.assertEqual(self.call_on_tx_unconfirmed(yg), ([], []))
+
+    def test_offer_changed(self):
+        yg, offer = self.create_yg_and_offer(100)
+        yg.offerlist = [{'oid': 0, 'maxsize': 10}]
+        self.assertEqual(self.call_on_tx_unconfirmed(yg), ([], [offer]))

scripts/yg-privacyenhanced.py

@@ -35,12 +35,10 @@ jlog = get_log()
 class YieldGeneratorPrivacyEnhanced(YieldGeneratorBasic):
 
     def __init__(self, wallet, offerconfig):
-        self.txfee, self.cjfee_a, self.cjfee_r, self.ordertype, self.minsize \
-            = offerconfig
         super(YieldGeneratorPrivacyEnhanced, self).__init__(wallet, offerconfig)
 
     def create_my_orders(self):
-        mix_balance = self.wallet.get_balance_by_mixdepth()
+        mix_balance = self.get_available_mixdepths(verbose=False)
         # We publish ONLY the maximum amount and use minsize for lower bound;
         # leave it to oid_to_order to figure out the right depth to use.
         f = '0'

scripts/yield-generator-basic.py

@@ -3,7 +3,7 @@ from __future__ import (absolute_import, division,
                         print_function, unicode_literals)
 from builtins import * # noqa: F401
 
-from jmbase import get_log, jmprint
+from jmbase import jmprint
 from jmclient import YieldGeneratorBasic, ygmain
 
 """THESE SETTINGS CAN SIMPLY BE EDITED BY HAND IN THIS FILE:
@@ -16,8 +16,6 @@ nickserv_password = ''
 max_minsize = 100000
 gaplimit = 6
 
-jlog = get_log()
-
 if __name__ == "__main__":
     ygmain(YieldGeneratorBasic, txfee=txfee, cjfee_a=cjfee_a,
            cjfee_r=cjfee_r, ordertype=ordertype,