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