src/allmydata/test/test_helper.py

   1
   2 import os
   3 from twisted.trial import unittest
   4 from twisted.application import service
   5
   6 from foolscap import Tub, eventual
   7 from foolscap.logging import log
   8
   9 from allmydata import upload, offloaded
  10 from allmydata.util import hashutil, fileutil
  11
  12 MiB = 1024*1024
  13
  14 class CHKUploadHelper_fake(offloaded.CHKUploadHelper):
  15     def start_encrypted(self, eu):
  16         d = eu.get_size()
  17         def _got_size(size):
  18             d2 = eu.get_all_encoding_parameters()
  19             def _got_parms(parms):
  20                 needed_shares, happy, total_shares, segsize = parms
  21                 return (hashutil.uri_extension_hash(""),
  22                         needed_shares, total_shares, size)
  23             d2.addCallback(_got_parms)
  24             return d2
  25         d.addCallback(_got_size)
  26         return d
  27
  28 class CHKUploadHelper_already_uploaded(offloaded.CHKUploadHelper):
  29     def start(self):
  30         res = {'uri_extension_hash': hashutil.uri_extension_hash("")}
  31         return (res, None)
  32
  33 class FakeClient(service.MultiService):
  34     DEFAULT_ENCODING_PARAMETERS = {"k":25,
  35                                    "happy": 75,
  36                                    "n": 100,
  37                                    "max_segment_size": 1*MiB,
  38                                    }
  39     def log(self, *args, **kwargs):
  40         return log.msg(*args, **kwargs)
  41     def get_push_to_ourselves(self):
  42         return True
  43     def get_encoding_parameters(self):
  44         return self.DEFAULT_ENCODING_PARAMETERS
  45
  46 def flush_but_dont_ignore(res):
  47     d = eventual.flushEventualQueue()
  48     def _done(ignored):
  49         return res
  50     d.addCallback(_done)
  51     return d
  52
  53 class AssistedUpload(unittest.TestCase):
  54     def setUp(self):
  55         self.s = FakeClient()
  56         self.s.startService()
  57
  58         self.tub = t = Tub()
  59         t.setServiceParent(self.s)
  60         self.s.tub = t
  61         # we never actually use this for network traffic, so it can use a
  62         # bogus host/port
  63         t.setLocation("bogus:1234")
  64
  65     def setUpHelper(self, basedir):
  66         fileutil.make_dirs(basedir)
  67         self.helper = h = offloaded.Helper(basedir)
  68         h.chk_upload_helper_class = CHKUploadHelper_fake
  69         h.setServiceParent(self.s)
  70         self.helper_furl = self.tub.registerReference(h)
  71
  72     def tearDown(self):
  73         d = self.s.stopService()
  74         d.addCallback(eventual.fireEventually)
  75         d.addBoth(flush_but_dont_ignore)
  76         return d
  77
  78
  79     def test_one(self):
  80         self.basedir = "helper/AssistedUpload/test_one"
  81         self.setUpHelper(self.basedir)
  82         u = upload.Uploader(self.helper_furl)
  83         u.setServiceParent(self.s)
  84
  85         # wait a few turns
  86         d = eventual.fireEventually()
  87         d.addCallback(eventual.fireEventually)
  88         d.addCallback(eventual.fireEventually)
  89
  90         def _ready(res):
  91             assert u._helper
  92
  93             DATA = "I need help\n" * 1000
  94             return u.upload_data(DATA)
  95         d.addCallback(_ready)
  96         def _uploaded(uri):
  97             assert "CHK" in uri
  98         d.addCallback(_uploaded)
  99
 100         def _check_empty(res):
 101             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 102             self.failUnlessEqual(files, [])
 103             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 104             self.failUnlessEqual(files, [])
 105         d.addCallback(_check_empty)
 106
 107         return d
 108
 109
 110     def test_already_uploaded(self):
 111         self.basedir = "helper/AssistedUpload/test_already_uploaded"
 112         self.setUpHelper(self.basedir)
 113         self.helper.chk_upload_helper_class = CHKUploadHelper_already_uploaded
 114         u = upload.Uploader(self.helper_furl)
 115         u.setServiceParent(self.s)
 116
 117         # wait a few turns
 118         d = eventual.fireEventually()
 119         d.addCallback(eventual.fireEventually)
 120         d.addCallback(eventual.fireEventually)
 121
 122         def _ready(res):
 123             assert u._helper
 124
 125             DATA = "I need help\n" * 1000
 126             return u.upload_data(DATA)
 127         d.addCallback(_ready)
 128         def _uploaded(uri):
 129             assert "CHK" in uri
 130         d.addCallback(_uploaded)
 131
 132         def _check_empty(res):
 133             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 134             self.failUnlessEqual(files, [])
 135             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 136             self.failUnlessEqual(files, [])
 137         d.addCallback(_check_empty)
 138
 139         return d