p2pool/test/util/test_forest.py

   1 import random
   2 import unittest
   3
   4 from p2pool.util import forest, math
   5 from p2pool.bitcoin import data as bitcoin_data
   6
   7 class DumbTracker(object):
   8     def __init__(self, shares=[]):
   9         self.shares = {} # hash -> share
  10         self.reverse_shares = {} # previous_hash -> set of share_hashes
  11
  12         for share in shares:
  13             self.add(share)
  14
  15     def add(self, share):
  16         if share.hash in self.shares:
  17             raise ValueError('share already present')
  18         self.shares[share.hash] = share
  19         self.reverse_shares.setdefault(share.previous_hash, set()).add(share.hash)
  20
  21     def remove(self, share_hash):
  22         share = self.shares[share_hash]
  23         del share_hash
  24
  25         self.shares.pop(share.hash)
  26         self.reverse_shares[share.previous_hash].remove(share.hash)
  27         if not self.reverse_shares[share.previous_hash]:
  28             self.reverse_shares.pop(share.previous_hash)
  29
  30     @property
  31     def heads(self):
  32         return dict((x, self.get_last(x)) for x in self.shares if x not in self.reverse_shares)
  33
  34     @property
  35     def tails(self):
  36         return dict((x, set(y for y in self.shares if self.get_last(y) == x and y not in self.reverse_shares)) for x in self.reverse_shares if x not in self.shares)
  37
  38     def get_height(self, share_hash):
  39         height, work, last = self.get_height_work_and_last(share_hash)
  40         return height
  41
  42     def get_work(self, share_hash):
  43         height, work, last = self.get_height_work_and_last(share_hash)
  44         return work
  45
  46     def get_last(self, share_hash):
  47         height, work, last = self.get_height_work_and_last(share_hash)
  48         return last
  49
  50     def get_height_and_last(self, share_hash):
  51         height, work, last = self.get_height_work_and_last(share_hash)
  52         return height, last
  53
  54     def get_nth_parent_hash(self, share_hash, n):
  55         for i in xrange(n):
  56             share_hash = self.shares[share_hash].previous_hash
  57         return share_hash
  58
  59     def get_height_work_and_last(self, share_hash):
  60         height = 0
  61         work = 0
  62         while share_hash in self.shares:
  63             share_hash, work_inc = self.shares[share_hash].previous_hash, bitcoin_data.target_to_average_attempts(self.shares[share_hash].target)
  64             height += 1
  65             work += work_inc
  66         return height, work, share_hash
  67
  68     def get_chain(self, start_hash, length):
  69         # same implementation :/
  70         assert length <= self.get_height(start_hash)
  71         for i in xrange(length):
  72             yield self.shares[start_hash]
  73             start_hash = self.shares[start_hash].previous_hash
  74
  75     def is_child_of(self, share_hash, possible_child_hash):
  76         if self.get_last(share_hash) != self.get_last(possible_child_hash):
  77             return None
  78         while True:
  79             if possible_child_hash == share_hash:
  80                 return True
  81             if possible_child_hash not in self.shares:
  82                 return False
  83             possible_child_hash = self.shares[possible_child_hash].previous_hash
  84
  85 class FakeShare(object):
  86     target = 2**256 - 1
  87     def __init__(self, **kwargs):
  88         for k, v in kwargs.iteritems():
  89             setattr(self, k, v)
  90         self._attrs = kwargs
  91
  92     def __repr__(self):
  93         return 'FakeShare(' + ', '.join('%s=%r' % (k, v) for k, v in self._attrs.iteritems()) + ')'
  94
  95 def test_tracker(self):
  96     t = DumbTracker(self.shares.itervalues())
  97
  98     assert self.shares == t.shares, (self.shares, t.shares)
  99     assert self.reverse_shares == t.reverse_shares, (self.reverse_shares, t.reverse_shares)
 100     assert self.heads == t.heads, (self.heads, t.heads)
 101     assert self.tails == t.tails, (self.tails, t.tails)
 102
 103     for start in self.shares:
 104         a, b = self.get_height_work_and_last(start), t.get_height_work_and_last(start)
 105         assert a == b, (a, b)
 106
 107         other = random.choice(self.shares.keys())
 108         assert self.is_child_of(start, other) == t.is_child_of(start, other)
 109         assert self.is_child_of(other, start) == t.is_child_of(other, start)
 110
 111         assert list(self.get_chain(start, min(a[0], 10))) == list(t.get_chain(start, min(a[0], 10)))
 112
 113 def generate_tracker_simple(n):
 114     t = forest.Tracker(math.shuffled(FakeShare(hash=i, previous_hash=i - 1 if i > 0 else None) for i in xrange(n)))
 115     test_tracker(t)
 116     return t
 117
 118 def generate_tracker_random(n):
 119     shares = []
 120     for i in xrange(n):
 121         x = random.choice(shares + [FakeShare(hash=None), FakeShare(hash=random.randrange(1000000, 2000000))]).hash
 122         shares.append(FakeShare(hash=i, previous_hash=x))
 123     t = forest.Tracker(math.shuffled(shares))
 124     test_tracker(t)
 125     return t
 126
 127 class Test(unittest.TestCase):
 128     def test_tracker(self):
 129         t = generate_tracker_simple(1000)
 130
 131         assert t.heads == {999: None}
 132         assert t.tails == {None: set([999])}
 133
 134         assert t.get_nth_parent_hash(900, 500) == 900 - 500
 135         assert t.get_nth_parent_hash(901, 42) == 901 - 42
 136
 137     def test_get_nth_parent_hash(self):
 138         t = generate_tracker_simple(200)
 139
 140         for i in xrange(1000):
 141             a = random.randrange(200)
 142             b = random.randrange(a + 1)
 143             res = t.get_nth_parent_hash(a, b)
 144             assert res == a - b, (a, b, res)
 145
 146     def test_tracker2(self):
 147         for ii in xrange(50):
 148             t = generate_tracker_random(random.randrange(100))
 149             #print "--start--"
 150             while t.shares:
 151                 x = random.choice(list(t.shares))
 152                 try:
 153                     t.remove(x)
 154                 except NotImplementedError:
 155                     pass # print "aborted", x
 156                 test_tracker(t)