tahoe-lafs/src/allmydata/test/test_helper.py

108 lines
3.0 KiB
Python

from twisted.trial import unittest
from twisted.application import service
from foolscap import Tub, eventual
from foolscap.logging import log
from allmydata import upload, offloaded
from allmydata.util import hashutil
class CHKUploadHelper_fake(offloaded.CHKUploadHelper):
def start_encrypted(self, eu):
needed_shares, happy, total_shares = self._encoding_parameters
d = eu.get_size()
def _got_size(size):
return (hashutil.uri_extension_hash(""),
needed_shares, total_shares, size)
d.addCallback(_got_size)
return d
class CHKUploadHelper_already_uploaded(offloaded.CHKUploadHelper):
def start(self):
res = {'uri_extension_hash': hashutil.uri_extension_hash("")}
return (res, None)
class FakeClient(service.MultiService):
def log(self, msg, **kwargs):
return log.msg(msg, **kwargs)
def get_push_to_ourselves(self):
return True
def get_encoding_parameters(self):
return None
def flush_but_dont_ignore(res):
d = eventual.flushEventualQueue()
def _done(ignored):
return res
d.addCallback(_done)
return d
class AssistedUpload(unittest.TestCase):
def setUp(self):
self.s = FakeClient()
self.s.startService()
self.tub = t = Tub()
t.setServiceParent(self.s)
self.s.tub = t
# we never actually use this for network traffic, so it can use a
# bogus host/port
t.setLocation("bogus:1234")
self.helper = h = offloaded.Helper(".")
h.chk_upload_helper_class = CHKUploadHelper_fake
h.setServiceParent(self.s)
self.helper_furl = t.registerReference(h)
def tearDown(self):
d = self.s.stopService()
d.addCallback(eventual.fireEventually)
d.addBoth(flush_but_dont_ignore)
return d
def test_one(self):
u = upload.Uploader(self.helper_furl)
u.setServiceParent(self.s)
# wait a few turns
d = eventual.fireEventually()
d.addCallback(eventual.fireEventually)
d.addCallback(eventual.fireEventually)
def _ready(res):
assert u._helper
DATA = "I need help\n" * 1000
return u.upload_data(DATA)
d.addCallback(_ready)
def _uploaded(uri):
assert "CHK" in uri
d.addCallback(_uploaded)
return d
def test_already_uploaded(self):
self.helper.chk_upload_helper_class = CHKUploadHelper_already_uploaded
u = upload.Uploader(self.helper_furl)
u.setServiceParent(self.s)
# wait a few turns
d = eventual.fireEventually()
d.addCallback(eventual.fireEventually)
d.addCallback(eventual.fireEventually)
def _ready(res):
assert u._helper
DATA = "I need help\n" * 1000
return u.upload_data(DATA)
d.addCallback(_ready)
def _uploaded(uri):
assert "CHK" in uri
d.addCallback(_uploaded)
return d