src/allmydata/test/test_helper.py

   1
   2 from twisted.trial import unittest
   3 from twisted.application import service
   4
   5 from foolscap import Tub, eventual
   6 from foolscap.logging import log
   7
   8 from allmydata import upload, offloaded
   9 from allmydata.util import hashutil
  10
  11 class CHKUploadHelper_fake(offloaded.CHKUploadHelper):
  12     def start_encrypted(self, eu):
  13         needed_shares, happy, total_shares = self._encoding_parameters
  14         d = eu.get_size()
  15         def _got_size(size):
  16             return (hashutil.uri_extension_hash(""),
  17                     needed_shares, total_shares, size)
  18         d.addCallback(_got_size)
  19         return d
  20
  21 class CHKUploadHelper_already_uploaded(offloaded.CHKUploadHelper):
  22     def start(self):
  23         res = {'uri_extension_hash': hashutil.uri_extension_hash("")}
  24         return (res, None)
  25
  26 class FakeClient(service.MultiService):
  27     def log(self, msg, **kwargs):
  28         return log.msg(msg, **kwargs)
  29     def get_push_to_ourselves(self):
  30         return True
  31     def get_encoding_parameters(self):
  32         return None
  33
  34 def flush_but_dont_ignore(res):
  35     d = eventual.flushEventualQueue()
  36     def _done(ignored):
  37         return res
  38     d.addCallback(_done)
  39     return d
  40
  41 class AssistedUpload(unittest.TestCase):
  42     def setUp(self):
  43         self.s = FakeClient()
  44         self.s.startService()
  45
  46         self.tub = t = Tub()
  47         t.setServiceParent(self.s)
  48         self.s.tub = t
  49         # we never actually use this for network traffic, so it can use a
  50         # bogus host/port
  51         t.setLocation("bogus:1234")
  52
  53         self.helper = h = offloaded.Helper(".")
  54         h.chk_upload_helper_class = CHKUploadHelper_fake
  55         h.setServiceParent(self.s)
  56         self.helper_furl = t.registerReference(h)
  57
  58     def tearDown(self):
  59         d = self.s.stopService()
  60         d.addCallback(eventual.fireEventually)
  61         d.addBoth(flush_but_dont_ignore)
  62         return d
  63
  64
  65     def test_one(self):
  66         u = upload.Uploader(self.helper_furl)
  67         u.setServiceParent(self.s)
  68
  69         # wait a few turns
  70         d = eventual.fireEventually()
  71         d.addCallback(eventual.fireEventually)
  72         d.addCallback(eventual.fireEventually)
  73
  74         def _ready(res):
  75             assert u._helper
  76
  77             DATA = "I need help\n" * 1000
  78             return u.upload_data(DATA)
  79         d.addCallback(_ready)
  80         def _uploaded(uri):
  81             assert "CHK" in uri
  82         d.addCallback(_uploaded)
  83
  84         return d
  85
  86
  87     def test_already_uploaded(self):
  88         self.helper.chk_upload_helper_class = CHKUploadHelper_already_uploaded
  89         u = upload.Uploader(self.helper_furl)
  90         u.setServiceParent(self.s)
  91
  92         # wait a few turns
  93         d = eventual.fireEventually()
  94         d.addCallback(eventual.fireEventually)
  95         d.addCallback(eventual.fireEventually)
  96
  97         def _ready(res):
  98             assert u._helper
  99
 100             DATA = "I need help\n" * 1000
 101             return u.upload_data(DATA)
 102         d.addCallback(_ready)
 103         def _uploaded(uri):
 104             assert "CHK" in uri
 105         d.addCallback(_uploaded)
 106
 107         return d
 108