Allow users to specify object sizes with k/M/G suffixes

Change-Id: Ie8671e2313b76376d123c9d4eafcc1aff44ba0b8

3 files changed, 59 insertions, 4 deletions

diff --git a/swiftbench/bench.py b/swiftbench/bench.py
index b7f35a1..976c03a 100644
--- a/swiftbench/bench.py
+++ b/swiftbench/bench.py
@@ -33,7 +33,8 @@ from six.moves import range
 
 import swiftclient as client
 
-from swiftbench.utils import config_true_value, using_http_proxy
+from swiftbench.utils import config_true_value, using_http_proxy, \
+    get_size_bytes
 
 
 try:
@@ -219,10 +220,10 @@ class Bench(object):
             logger.warn("Communication with Swift server is going through "
                         "HTTP proxy server. This could affect test result")
 
-        self.object_size = int(conf.object_size)
+        self.object_size = get_size_bytes(conf.object_size)
         self.object_sources = conf.object_sources
-        self.lower_object_size = int(conf.lower_object_size)
-        self.upper_object_size = int(conf.upper_object_size)
+        self.lower_object_size = get_size_bytes(conf.lower_object_size)
+        self.upper_object_size = get_size_bytes(conf.upper_object_size)
         self.files = []
         if self.object_sources:
             self.object_sources = self.object_sources.split()
diff --git a/swiftbench/utils.py b/swiftbench/utils.py
index df39bab..07c4425 100644
--- a/swiftbench/utils.py
+++ b/swiftbench/utils.py
@@ -89,3 +89,22 @@ def using_http_proxy(url):
     """
     up = urlparse(url)
     return up.scheme.lower() in getproxies() and not proxy_bypass(up.netloc)
+
+
+def get_size_bytes(value):
+    """
+    Expand suffixes for object sizes.
+
+    For example, '10k' becomes 10240, and '2M' becomes 2097152.
+    """
+    if not isinstance(value, six.string_types):
+        raise TypeError
+    value = value.strip()
+    multiple = {
+        'G': 1024 ** 3,
+        'M': 1024 ** 2,
+        'k': 1024,
+    }.get(value[-1])
+    if multiple is None:
+        return int(value)
+    return multiple * int(value[:-1])
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 9be68b5..5859027 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -120,5 +120,40 @@ log_name = %(yarr)s'''
         self.assertFalse(utils.using_http_proxy("dummy://localhost/"))
         self.assertFalse(utils.using_http_proxy("dummy://host1/"))
 
+    def test_get_size_bytes(self):
+        self.assertEqual(utils.get_size_bytes('10M'), 10 * 1024 * 1024)
+        self.assertEqual(utils.get_size_bytes('4k'), 4096)
+        self.assertEqual(utils.get_size_bytes('5G'), 5 * 1024 * 1024 * 1024)
+        self.assertEqual(utils.get_size_bytes('1234'), 1234)
+        self.assertEqual(utils.get_size_bytes('1 k'), 1024)
+        self.assertEqual(utils.get_size_bytes(' 1k'), 1024)
+        self.assertEqual(utils.get_size_bytes('1k '), 1024)
+
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1K')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1m')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1g')
+
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1kb')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1kB')
+
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1T')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1P')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1E')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('1Y')
+        with self.assertRaises(ValueError):
+            utils.get_size_bytes('asdf')
+
+        with self.assertRaises(TypeError):
+            utils.get_size_bytes(1)
+
 if __name__ == '__main__':
     unittest.main()