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