allmydata/test/test_upload.py

   1
   2 from twisted.trial import unittest
   3 from twisted.internet import defer
   4 from cStringIO import StringIO
   5
   6 from allmydata import upload
   7
   8 class StringBucketProxy:
   9     # This is for unit tests: make a StringIO look like a RIBucketWriter.
  10
  11     def __init__(self):
  12         self.data = StringIO()
  13         self.size = None
  14         self.done = False
  15
  16     def callRemote(self, methname, **kwargs):
  17         if methname == "write":
  18             return defer.maybeDeferred(self.write, **kwargs)
  19         elif methname == "close":
  20             return defer.maybeDeferred(self.close, **kwargs)
  21         else:
  22             return defer.fail(NameError("no such method named %s" % methname))
  23
  24     def write(self, data):
  25         self.data.write(data)
  26     def close(self):
  27         self.done = True
  28
  29
  30 class FakePeer:
  31     def __init__(self, peerid, response):
  32         self.peerid = peerid
  33         self.response = response
  34
  35     def callRemote(self, methname, *args, **kwargs):
  36         assert not args
  37         return defer.maybeDeferred(self._callRemote, methname, **kwargs)
  38
  39     def _callRemote(self, methname, **kwargs):
  40         assert methname == "allocate_bucket"
  41         assert kwargs["size"] == 100
  42         assert kwargs["leaser"] == "fakeclient"
  43         if self.response == "good":
  44             return self
  45         raise upload.TooFullError()
  46
  47 class FakeClient:
  48     nodeid = "fakeclient"
  49     def __init__(self, responses):
  50         self.peers = []
  51         for peerid,r in enumerate(responses):
  52             if r == "disconnected":
  53                 self.peers.append(None)
  54             else:
  55                 self.peers.append(FakePeer(str(peerid), r))
  56
  57     def permute_peerids(self, key, max_peers):
  58         assert max_peers == None
  59         return [str(i) for i in range(len(self.peers))]
  60
  61     def get_remote_service(self, peerid, name):
  62         peer = self.peers[int(peerid)]
  63         if not peer:
  64             return defer.fail(IndexError("no connection to that peer"))
  65         return defer.succeed(peer)
  66
  67 class NextPeerUploader(upload.FileUploader):
  68     def _got_all_peers(self, res):
  69         return res
  70
  71 class NextPeer(unittest.TestCase):
  72     responses = ["good", # 0
  73                  "full", # 1
  74                  "full", # 2
  75                  "disconnected", # 3
  76                  "good", # 4
  77                  ]
  78
  79     def compare_landlords(self, u, c, expected):
  80         exp = [(str(peerid), bucketnum, c.peers[peerid])
  81                for peerid, bucketnum in expected]
  82         self.failUnlessEqual(u.landlords, exp)
  83
  84     def test_0(self):
  85         c = FakeClient([])
  86         u = NextPeerUploader(c)
  87         u._verifierid = "verifierid"
  88         u._shares = 2
  89         u._share_size = 100
  90         d = u.start()
  91         def _check(f):
  92             f.trap(upload.NotEnoughPeersError)
  93         d.addCallbacks(lambda res: self.fail("this was supposed to fail"),
  94                        _check)
  95         return d
  96
  97     def test_1(self):
  98         c = FakeClient(self.responses)
  99         u = NextPeerUploader(c)
 100         u._verifierid = "verifierid"
 101         u._shares = 2
 102         u._share_size = 100
 103         d = u.start()
 104         def _check(res):
 105             self.failUnlessEqual(u.goodness_points, 2)
 106             self.compare_landlords(u, c, [(0, 0),
 107                                           (4, 1),
 108                                           ])
 109         d.addCallback(_check)
 110         return d
 111
 112     def test_2(self):
 113         c = FakeClient(self.responses)
 114         u = NextPeerUploader(c)
 115         u._verifierid = "verifierid"
 116         u._shares = 3
 117         u._share_size = 100
 118         d = u.start()
 119         def _check(res):
 120             self.failUnlessEqual(u.goodness_points, 3)
 121             self.compare_landlords(u, c, [(0, 0),
 122                                           (4, 1),
 123                                           (0, 2),
 124                                           ])
 125         d.addCallback(_check)
 126         return d
 127
 128     responses2 = ["good", # 0
 129                  "full", # 1
 130                  "full", # 2
 131                  "good", # 3
 132                  "full", # 4
 133                  ]
 134
 135     def test_3(self):
 136         c = FakeClient(self.responses2)
 137         u = NextPeerUploader(c)
 138         u._verifierid = "verifierid"
 139         u._shares = 3
 140         u._share_size = 100
 141         d = u.start()
 142         def _check(res):
 143             self.failUnlessEqual(u.goodness_points, 3)
 144             self.compare_landlords(u, c, [(0, 0),
 145                                           (3, 1),
 146                                           (0, 2),
 147                                           ])
 148         d.addCallback(_check)
 149         return d
 150
 151     responses3 = ["good", # 0
 152                  "good", # 1
 153                  "good", # 2
 154                  "good", # 3
 155                  "good", # 4
 156                  ]
 157
 158     def test_4(self):
 159         c = FakeClient(self.responses3)
 160         u = NextPeerUploader(c)
 161         u._verifierid = "verifierid"
 162         u._shares = 4
 163         u._share_size = 100
 164         d = u.start()
 165         def _check(res):
 166             self.failUnlessEqual(u.goodness_points, 4)
 167             self.compare_landlords(u, c, [(0, 0),
 168                                           (1, 1),
 169                                           (2, 2),
 170                                           (3, 3),
 171                                           ])
 172         d.addCallback(_check)
 173         return d