tahoe-lafs/allmydata/test/test_upload.py

154 lines
4.6 KiB
Python

from twisted.trial import unittest
from twisted.internet import defer
from cStringIO import StringIO
from allmydata import upload
class StringBucketProxy:
# This is for unit tests: make a StringIO look like a RIBucketWriter.
def __init__(self):
self.data = StringIO()
self.size = None
self.done = False
def callRemote(self, methname, **kwargs):
if methname == "write":
return defer.maybeDeferred(self.write, **kwargs)
elif methname == "set_size":
return defer.maybeDeferred(self.set_size, **kwargs)
elif methname == "close":
return defer.maybeDeferred(self.close, **kwargs)
else:
return defer.fail(NameError("no such method named %s" % methname))
def write(self, data):
self.data.write(data)
def set_size(self, size):
self.size = size
def close(self):
self.done = True
class FakePeer:
def __init__(self, peerid, response):
self.peerid = peerid
self.response = response
def callRemote(self, methname, *args, **kwargs):
assert not args
return defer.maybeDeferred(self._callRemote, methname, **kwargs)
def _callRemote(self, methname, **kwargs):
assert methname == "allocate_bucket"
assert kwargs["size"] == 100
assert kwargs["leaser"] == "fakeclient"
if self.response == "good":
return self
raise upload.TooFullError()
class FakeClient:
nodeid = "fakeclient"
def __init__(self, responses):
self.peers = []
for peerid,r in enumerate(responses):
if r == "disconnected":
self.peers.append(None)
else:
self.peers.append(FakePeer(str(peerid), r))
def permute_peerids(self, key, max_peers):
assert max_peers == None
return [str(i) for i in range(len(self.peers))]
def get_remote_service(self, peerid, name):
peer = self.peers[int(peerid)]
if not peer:
return defer.fail(IndexError("no connection to that peer"))
return defer.succeed(peer)
class NextPeerUploader(upload.FileUploader):
def _got_all_peers(self, res):
return res
class NextPeer(unittest.TestCase):
responses = ["good", # 0
"full", # 1
"full", # 2
"disconnected", # 3
"good", # 4
]
def compare_landlords(self, u, c, expected):
exp = [(str(peerid), bucketnum, c.peers[peerid])
for peerid, bucketnum in expected]
self.failUnlessEqual(u.landlords, exp)
def test_0(self):
c = FakeClient([])
u = NextPeerUploader(c)
u._verifierid = "verifierid"
u._shares = 2
u._share_size = 100
d = u.start()
def _check(f):
f.trap(upload.NotEnoughPeersError)
d.addCallbacks(lambda res: self.fail("this was supposed to fail"),
_check)
return d
def test_1(self):
c = FakeClient(self.responses)
u = NextPeerUploader(c)
u._verifierid = "verifierid"
u._shares = 2
u._share_size = 100
d = u.start()
def _check(res):
self.failUnlessEqual(u.goodness_points, 2)
self.compare_landlords(u, c, [(0, 0),
(4, 1),
])
d.addCallback(_check)
return d
def test_2(self):
c = FakeClient(self.responses)
u = NextPeerUploader(c)
u._verifierid = "verifierid"
u._shares = 3
u._share_size = 100
d = u.start()
def _check(res):
self.failUnlessEqual(u.goodness_points, 3)
self.compare_landlords(u, c, [(0, 0),
(4, 1),
(0, 2),
])
d.addCallback(_check)
return d
responses2 = ["good", # 0
"full", # 1
"full", # 2
"good", # 3
"full", # 4
]
def test_3(self):
c = FakeClient(self.responses2)
u = NextPeerUploader(c)
u._verifierid = "verifierid"
u._shares = 3
u._share_size = 100
d = u.start()
def _check(res):
self.failUnlessEqual(u.goodness_points, 3)
self.compare_landlords(u, c, [(0, 0),
(3, 1),
(0, 2),
])
d.addCallback(_check)
return d