summaryrefslogtreecommitdiff
path: root/scripts/scalability/test_controller/client/templates.py
diff options
context:
space:
mode:
authordrebs <drebs@riseup.net>2017-10-20 10:42:52 -0200
committerdrebs <drebs@riseup.net>2017-11-02 09:36:31 -0200
commitbba22c8439b56df4a74583f2582a89045ee43182 (patch)
tree5569f481350349e0d3141836a1afa5fce79e17ab /scripts/scalability/test_controller/client/templates.py
parent49e7efc43f015bf9cf2de4ae1b5c632d57d04970 (diff)
[benchmarks] add blob template creator for scalability test
Diffstat (limited to 'scripts/scalability/test_controller/client/templates.py')
-rwxr-xr-xscripts/scalability/test_controller/client/templates.py89
1 files changed, 89 insertions, 0 deletions
diff --git a/scripts/scalability/test_controller/client/templates.py b/scripts/scalability/test_controller/client/templates.py
new file mode 100755
index 00000000..dcf51fe1
--- /dev/null
+++ b/scripts/scalability/test_controller/client/templates.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python
+
+import base64
+import errno
+import mock
+import os
+import random
+
+from argparse import ArgumentParser
+from io import BytesIO
+from tempfile import mkdtemp
+
+from twisted.internet import reactor, defer
+
+from leap.soledad.client._db.blobs import BlobManager
+
+
+DEFAULT_TARGET_DIR = './blob-templates'
+
+
+def _get_put_function(path):
+ def _put(_, data, params=None):
+ with open(path, 'w') as f:
+ f.write(data.read())
+ return defer.succeed(mock.Mock(code=200))
+ return _put
+
+
+def _get_encrypt_function(user, path):
+ tempdir = mkdtemp()
+ manager = BlobManager(tempdir, None, '123', '123', user)
+ manager._client.put = _get_put_function(path)
+ return manager._encrypt_and_upload
+
+
+def payload(size):
+ random.seed(1337) # same seed to avoid different bench results
+ payload_bytes = bytearray(random.getrandbits(8) for _ in xrange(size))
+ # encode as base64 to avoid ascii encode/decode errors
+ return base64.b64encode(payload_bytes)[:size] # remove b64 overhead
+
+
+def _encrypt(path, data):
+ encrypt = _get_encrypt_function('user-0', path)
+ return encrypt('blob', BytesIO(data))
+
+
+def mkdir_p(path):
+ try:
+ os.makedirs(path)
+ except OSError as exc: # Python >2.5
+ if exc.errno == errno.EEXIST and os.path.isdir(path):
+ pass
+ else:
+ raise
+
+
+def create_blob_templates(target_dir, amount, size):
+ mkdir_p(target_dir)
+ data = payload(size)
+ semaphore = defer.DeferredSemaphore(20)
+ deferreds = []
+ for i in xrange(amount):
+ path = os.path.join(target_dir,
+ '%dK-%d.blob' % (size / 1000, i))
+ d = semaphore.run(_encrypt, path, data)
+ deferreds.append(d)
+ return defer.gatherResults(deferreds)
+
+
+def parse_args():
+ parser = ArgumentParser()
+ parser.add_argument(
+ 'target_dir',
+ help='The target directory where templates will be written to.')
+ parser.add_argument(
+ '--amount', default=1000, type=int,
+ help='The number of blob templates to create.')
+ parser.add_argument(
+ '--size', default=1000, type=int,
+ help='The size of each template in KB.')
+ return parser.parse_args()
+
+
+if __name__ == '__main__':
+ args = parse_args()
+ d = create_blob_templates(args.target_dir, args.amount, args.size * 1000)
+ d.addCallback(lambda _: reactor.stop())
+ reactor.run()