From b87114e40c5f095103be4a9339b38552333d0190 Mon Sep 17 00:00:00 2001
From: Pete Vander Giessen <petevg@gmail.com>
Date: Mon, 28 Nov 2016 13:26:51 -0500
Subject: [PATCH] Fixed addMachines.

We now accept a single dict called params, which is what we get back
from the planner, rather than expecting expanded arguments.

We do some parsing of constraints, to put them in a format that should
make juju core happy.
---
 juju/constraints.py            | 64 ++++++++++++++++++++++++++++++++++
 juju/model.py                  | 28 +++++++++------
 tests/unit/test_constraints.py | 45 ++++++++++++++++++++++++
 3 files changed, 126 insertions(+), 11 deletions(-)
 create mode 100644 juju/constraints.py
 create mode 100644 tests/unit/test_constraints.py

diff --git a/juju/constraints.py b/juju/constraints.py
new file mode 100644
index 0000000..d713774
--- /dev/null
+++ b/juju/constraints.py
@@ -0,0 +1,64 @@
+#
+# Library to parse constraints
+#
+# The current version of juju core expects the client to take
+# constraints given in the form "mem=10G foo=bar" and parse them into
+# json that looks like {"mem": 10240, "foo": "bar"}. This module helps us
+# accomplish that task.
+#
+# We do not attempt to duplicate the checking done in
+# client/_client.py:Value here. That class will verify that the
+# constraints keys are valid, and that we can successfully dump the
+# constraints dict to json.
+#
+
+import re
+
+# Matches on a string specifying memory size
+MEM = re.compile('^[1-9][0-9]*[MGTP]$')
+
+# Multiplication factors to get Megabytes
+FACTORS = {
+    "M": 1,
+    "G": 1024,
+    "T": 1024 * 1024,
+    "P": 1024 * 1024 * 1024
+}
+
+def parse(constraints):
+    """
+    Constraints must be expressed as a string containing only spaces
+    and key value pairs joined by an '='.
+
+    """
+    if constraints is None:
+        return None
+
+    constraints = {
+        normalize_key(k): normalize_value(v) for k, v in [
+            s.split("=") for s in constraints.split(" ")]}
+
+    return constraints
+
+
+def normalize_key(key):
+    key = key.strip()
+
+    key = key.replace("-", "_")  # Our _client lib wants "_" in place of "-"
+    return key
+
+
+def normalize_value(value):
+    value = value.strip()
+
+    if MEM.match(value):
+        # Translate aliases to Megabytes. e.g. 1G = 10240
+        return int(value[:-1]) * FACTORS[value[-1:]]
+
+    if "," in value:
+        # Handle csv strings.
+        values = value.split(",")
+        values = [normalize_value(v) for v in values]
+        return values
+
+    return value
diff --git a/juju/model.py b/juju/model.py
index 9d14f82..e412784 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -12,6 +12,7 @@ from theblues import charmstore
 from .client import client
 from .client import watcher
 from .client import connection
+from .constraints import parse as parse_constraints
 from .delta import get_entity_delta
 from .delta import get_entity_class
 from .exceptions import DeadEntityException
@@ -1255,8 +1256,7 @@ class BundleHandler(object):
         await self.client_facade.AddCharm(None, entity_id)
         return entity_id
 
-    async def addMachines(self, series, constraints, container_type,
-                          parent_id):
+    async def addMachines(self, params):
         """
         :param series string:
             Series holds the optional machine OS series.
@@ -1275,15 +1275,21 @@ class BundleHandler(object):
             case this machine is a container, in which case also ContainerType
             is set.
         """
-        params = client.AddMachineParams(
-            series=series,
-            constraints=constraints,
-            container_type=container_type,
-            parent_id=self.resolve(parent_id),
-        )
-        results = await self.client_facade.AddMachines(params)
-        log.debug('Added new machine %s', results[0].machine)
-        return results[0].machine
+        if 'parent_id' in params:
+            params['parent_id'] = self.resolve(params['parent_id'])
+
+        params['constraints'] = parse_constraints(
+            params.get('constraints'))
+        params['jobs'] = params.get('jobs', ['JobHostUnits'])
+
+        params = client.AddMachineParams(**params)
+        results = await self.client_facade.AddMachines([params])
+        error = results.machines[0].error
+        if error:
+            raise ValueError("Error adding machine: %s", error.message)
+        machine = results.machines[0].machine
+        log.debug('Added new machine %s', machine)
+        return machine
 
     async def addRelation(self, endpoint1, endpoint2):
         """
diff --git a/tests/unit/test_constraints.py b/tests/unit/test_constraints.py
new file mode 100644
index 0000000..ec50bdd
--- /dev/null
+++ b/tests/unit/test_constraints.py
@@ -0,0 +1,45 @@
+#
+# Test our constraints parser
+#
+
+import unittest
+
+from juju import constraints
+
+class TestConstraints(unittest.TestCase):
+
+    def test_mem_regex(self):
+        m = constraints.MEM
+        self.assertTrue(m.match("10G"))
+        self.assertTrue(m.match("1G"))
+        self.assertFalse(m.match("1Gb"))
+        self.assertFalse(m.match("a1G"))
+        self.assertFalse(m.match("1000"))
+
+    def test_normalize_key(self):
+        _ = constraints.normalize_key
+
+        self.assertEqual(_("test-key"), "test_key")
+        self.assertEqual(_("test-key  "), "test_key")
+        self.assertEqual(_("  test-key"), "test_key")
+
+    def test_normalize_val(self):
+        _ = constraints.normalize_value
+
+        self.assertEqual(_("10G"), 10 * 1024)
+        self.assertEqual(_("10M"), 10)
+        self.assertEqual(_("10"), "10")
+        self.assertEqual(_("foo,bar"), ["foo", "bar"])
+
+    def test_parse_constraints(self):
+        _ = constraints.parse
+
+        self.assertEqual(
+            _("mem=10G"),
+            {"mem": 10 * 1024}
+        )
+
+        self.assertEqual(
+            _("mem=10G foo=bar,baz"),
+            {"mem": 10 * 1024, "foo": ["bar", "baz"]}
+        )
-- 
2.25.1