Merge lp:~danilo/maas/virsh-storage-units-2.2 into lp:maas/2.2

[Данило Шеган]

Backport r6095 from trunk.

Consider storage size units when figuring out virsh pod storage pool sizes.

Backport of https://code.launchpad.net/~danilo/maas/virsh-storage-units-trunk/+merge/325920

[Данило Шеган (danilo)]

Self-approving a backport.

review: Approve

=== modified file 'src/provisioningserver/drivers/pod/tests/test_virsh.py'
--- src/provisioningserver/drivers/pod/tests/test_virsh.py	2017-05-19 03:37:56 +0000
+++ src/provisioningserver/drivers/pod/tests/test_virsh.py	2017-06-21 11:02:31 +0000
@@ -105,6 +105,17 @@
     Available:      173.84 GiB
     """)
 
+SAMPLE_POOLINFO_TB = dedent("""
+    Name:           default
+    UUID:           59edc0cb-4635-449a-80e2-2c8a59afa327
+    State:          running
+    Persistent:     yes
+    Autostart:      yes
+    Capacity:       2.21 TiB
+    Allocation:     880.64 GiB
+    Available:      1.35 TiB
+    """)
+
 SAMPLE_IFLIST = dedent("""
     Interface  Type       Source     Model       MAC
     -------------------------------------------------------
@@ -506,6 +517,18 @@
             for pool in pools_mock.return_value
         }, expected)
 
+    def test_get_pod_pool_size_map_terabytes(self):
+        conn = self.configure_virshssh(SAMPLE_POOLINFO_TB)
+        pools_mock = self.patch(virsh.VirshSSH, 'list_pools')
+        pools_mock.return_value = [
+            factory.make_name('pool') for _ in range(3)]
+        expected = conn.get_pod_pool_size_map('Capacity')
+        capacity = int(2.21 * 2**40)
+        self.assertEqual({
+            pool: capacity
+            for pool in pools_mock.return_value
+        }, expected)
+
     def test_get_pod_local_storage(self):
         conn = self.configure_virshssh(SAMPLE_POOLINFO)
         pools_mock = self.patch(virsh.VirshSSH, 'list_pools')
 
=== modified file 'src/provisioningserver/drivers/pod/virsh.py'
--- src/provisioningserver/drivers/pod/virsh.py	2017-05-19 03:57:03 +0000
+++ src/provisioningserver/drivers/pod/virsh.py	2017-06-21 11:02:31 +0000
@@ -37,6 +37,7 @@
     create_node,
 )
 from provisioningserver.utils import (
+    convert_size_to_bytes,
     shell,
     typed,
 )
@@ -273,7 +274,13 @@
         data = data.strip().splitlines()
         for d in data:
             if key == d.split(':')[0].strip():
-                return d.split(':')[1].split()[0]
+                return d.split(':')[1].strip()
+
+    def get_key_value_unitless(self, data, key):
+        """Return value based off of key with unit (if any) stripped off."""
+        value = self.get_key_value(data, key)
+        if value:
+            return value.split()[0]
 
     def list_machines(self):
         """Lists all VMs by name."""
@@ -336,7 +343,7 @@
         if output is None:
             maaslog.error("Failed to get pod CPU speed")
             return None
-        return int(self.get_key_value(output, "CPU frequency"))
+        return int(self.get_key_value_unitless(output, "CPU frequency"))
 
     def get_pod_memory(self):
         """Gets the total memory of the pod."""
@@ -344,7 +351,7 @@
         if output is None:
             maaslog.error("Failed to get pod memory")
             return None
-        KiB = int(self.get_key_value(output, "Memory size"))
+        KiB = int(self.get_key_value_unitless(output, "Memory size"))
         # Memory in MiB.
         return int(KiB / 1024)
 
@@ -354,7 +361,7 @@
         if output is None:
             maaslog.error("%s: Failed to get machine memory", machine)
             return None
-        KiB = int(self.get_key_value(output, "Max memory"))
+        KiB = int(self.get_key_value_unitless(output, "Max memory"))
         # Memory in MiB.
         return int(KiB / 1024)
 
@@ -366,10 +373,8 @@
             if output is None:
                 # Skip if cannot get more information.
                 continue
-            capacity = float(
-                self.get_key_value(output, key))
-            # Convert GiB to bytes.
-            pools[pool] = int(capacity * 2**30)
+            pools[pool] = convert_size_to_bytes(
+                self.get_key_value(output, "Capacity"))
         return pools
 
     def get_pod_local_storage(self):
@@ -390,10 +395,10 @@
                 maaslog.error(
                     "Failed to get available pod local storage")
                 return None
-            local_storage += float(self.get_key_value(
-                output, "Available"))
-        # Local storage in bytes. GiB to bytes.
-        return int(local_storage * 2**30)
+            local_storage += convert_size_to_bytes(
+                self.get_key_value(output, "Available"))
+        # Local storage in bytes.
+        return local_storage
 
     def get_machine_local_storage(self, machine, device):
         """Gets the VM local storage for device."""
 
=== modified file 'src/provisioningserver/utils/__init__.py'
--- src/provisioningserver/utils/__init__.py	2017-04-05 14:57:42 +0000
+++ src/provisioningserver/utils/__init__.py	2017-06-21 11:02:31 +0000
@@ -246,3 +246,48 @@
         return issubclass(test, query_tuple)
     except TypeError:
         return False
+
+
+# Capacity units supported by convert_size_to_bytes() function.
+CAPACITY_UNITS = {
+    "KiB": 2 ** 10,
+    "MiB": 2 ** 20,
+    "GiB": 2 ** 30,
+    "TiB": 2 ** 40,
+    "PiB": 2 ** 50,
+    "EiB": 2 ** 60,
+    "ZiB": 2 ** 70,
+    "YiB": 2 ** 80,
+}
+
+
+class UnknownCapacityUnitError(Exception):
+    """Unknown capacity unit used."""
+
+
+def convert_size_to_bytes(value):
+    """
+    Converts storage size values with units (GiB, TiB...) to bytes.
+
+    :param value: A string containing a number and unit separated by at least
+        one space character.  If unit is not specified, defaults to bytes.
+    :return: An integer indicating the number of bytes for the given value in
+        any other size unit.
+    :raises UnknownCapacityUnitError: unsupported capacity unit.
+    """
+    # Split value on the first space.
+    capacity_def = value.split(" ", 1)
+    if len(capacity_def) == 1:
+        # No unit specified, default to bytes.
+        return int(capacity_def[0])
+
+    capacity_value, capacity_unit = capacity_def
+    capacity_value = float(capacity_value)
+    capacity_unit = capacity_unit.strip()
+    if capacity_unit in CAPACITY_UNITS:
+        multiplier = CAPACITY_UNITS[capacity_unit]
+    else:
+        raise UnknownCapacityUnitError(
+            "Unknown capacity unit '%s'" % capacity_unit)
+    # Convert value to bytes.
+    return int(capacity_value * multiplier)
 
=== modified file 'src/provisioningserver/utils/tests/test_utils.py'
--- src/provisioningserver/utils/tests/test_utils.py	2017-03-31 20:27:52 +0000
+++ src/provisioningserver/utils/tests/test_utils.py	2017-06-21 11:02:31 +0000
@@ -21,6 +21,7 @@
 from provisioningserver.utils import (
     CircularDependency,
     classify,
+    convert_size_to_bytes,
     flatten,
     is_instance_or_subclass,
     locate_config,
@@ -30,6 +31,7 @@
     ShellTemplate,
     sorttop,
     sudo,
+    UnknownCapacityUnitError,
 )
 from testtools.matchers import (
     DirExists,
@@ -369,3 +371,39 @@
         self.assertThat(
             is_instance_or_subclass(
                 self.bar, *[Baz, [Bar, [Foo]]]), Equals(True))
+
+
+class TestConvertSizeToBytes(MAASTestCase):
+    """Tests for `convert_size_to_bytes`."""
+
+    scenarios = (
+        ("bytes", {"value": "24111", "expected": 24111}),
+        ("KiB", {"value": "2.21 KiB", "expected": int(2.21 * 2**10)}),
+        ("MiB", {"value": "2.21 MiB", "expected": int(2.21 * 2**20)}),
+        ("GiB", {"value": "2.21 GiB", "expected": int(2.21 * 2**30)}),
+        ("TiB", {"value": "2.21 TiB", "expected": int(2.21 * 2**40)}),
+        ("PiB", {"value": "2.21 PiB", "expected": int(2.21 * 2**50)}),
+        ("EiB", {"value": "2.21 EiB", "expected": int(2.21 * 2**60)}),
+        ("ZiB", {"value": "2.21 ZiB", "expected": int(2.21 * 2**70)}),
+        ("YiB", {"value": "2.21 YiB", "expected": int(2.21 * 2**80)}),
+        ("whitespace", {"value": "2.21   GiB", "expected": int(2.21 * 2**30)}),
+        ("zero", {"value": "0 TiB", "expected": 0}),
+    )
+
+    def test__convert_size_to_bytes(self):
+        self.assertEqual(self.expected, convert_size_to_bytes(self.value))
+
+
+class TestConvertSizeToBytesErrors(MAASTestCase):
+    """Error handling tests for `convert_size_to_bytes`."""
+
+    def test__unknown_capacity_unit(self):
+        error = self.assertRaises(
+            UnknownCapacityUnitError, convert_size_to_bytes, "200 superbytes")
+        self.assertEqual("Unknown capacity unit 'superbytes'", str(error))
+
+    def test__empty_string(self):
+        self.assertRaises(ValueError, convert_size_to_bytes, "")
+
+    def test__empty_value(self):
+        self.assertRaises(ValueError, convert_size_to_bytes, " KiB")