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 is_child_of(self, share_hash, possible_child_hash):
  69         if self.get_last(share_hash) != self.get_last(possible_child_hash):
  70             return None
  71         while True:
  72             if possible_child_hash == share_hash:
  73                 return True
  74             if possible_child_hash not in self.shares:
  75                 return False
  76             possible_child_hash = self.shares[possible_child_hash].previous_hash
  77
  78 class FakeShare(object):
  79     target = 2**256 - 1
  80     def __init__(self, **kwargs):
  81         for k, v in kwargs.iteritems():
  82             setattr(self, k, v)
  83         self._attrs = kwargs
  84
  85     def __repr__(self):
  86         return 'FakeShare(' + ', '.join('%s=%r' % (k, v) for k, v in self._attrs.iteritems()) + ')'
  87
  88 def test_tracker(self):
  89     t = DumbTracker(self.shares.itervalues())
  90
  91     assert self.shares == t.shares, (self.shares, t.shares)
  92     assert self.reverse_shares == t.reverse_shares, (self.reverse_shares, t.reverse_shares)
  93     assert self.heads == t.heads, (self.heads, t.heads)
  94     assert self.tails == t.tails, (self.tails, t.tails)
  95
  96     for start in self.shares:
  97         a, b = self.get_height_work_and_last(start), t.get_height_work_and_last(start)
  98         assert a == b, (a, b)
  99
 100         other = random.choice(self.shares.keys())
 101         assert self.is_child_of(start, other) == t.is_child_of(start, other)
 102         assert self.is_child_of(other, start) == t.is_child_of(other, start)
 103
 104 def generate_tracker_simple(n):
 105     t = forest.Tracker(math.shuffled(FakeShare(hash=i, previous_hash=i - 1 if i > 0 else None) for i in xrange(n)))
 106     test_tracker(t)
 107     return t
 108
 109 def generate_tracker_random(n):
 110     shares = []
 111     for i in xrange(n):
 112         x = random.choice(shares + [FakeShare(hash=None), FakeShare(hash=random.randrange(1000000, 2000000))]).hash
 113         shares.append(FakeShare(hash=i, previous_hash=x))
 114     t = forest.Tracker(math.shuffled(shares))
 115     test_tracker(t)
 116     return t
 117
 118 class Test(unittest.TestCase):
 119     def test_tracker(self):
 120         t = generate_tracker_simple(1000)
 121
 122         assert t.heads == {999: None}
 123         assert t.tails == {None: set([999])}
 124
 125         assert t.get_nth_parent_hash(900, 500) == 900 - 500
 126         assert t.get_nth_parent_hash(901, 42) == 901 - 42
 127
 128     def test_get_nth_parent_hash(self):
 129         t = generate_tracker_simple(200)
 130
 131         for i in xrange(1000):
 132             a = random.randrange(200)
 133             b = random.randrange(a + 1)
 134             res = t.get_nth_parent_hash(a, b)
 135             assert res == a - b, (a, b, res)
 136
 137     def test_tracker2(self):
 138         for ii in xrange(50):
 139             t = generate_tracker_random(random.randrange(100))
 140             #print "--start--"
 141             while t.shares:
 142                 x = random.choice(list(t.shares))
 143                 try:
 144                     t.remove(x)
 145                 except NotImplementedError:
 146                     pass # print "aborted", x
 147                 test_tracker(t)