From 436db3acad1ea4d5221b556073405c7cbd70b425 Mon Sep 17 00:00:00 2001
From: Tim Van Steenburgh <tvansteenburgh@gmail.com>
Date: Fri, 10 Jun 2016 12:03:07 -0400
Subject: [PATCH] Add add_storage() method to UnitSentry

---
 amulet/sentry.py      | 13 +++++++++++++
 amulet/storage.py     | 30 ++++++++++++++++++++++++++++++
 tests/test_storage.py | 39 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 amulet/storage.py
 create mode 100644 tests/test_storage.py

diff --git a/amulet/sentry.py b/amulet/sentry.py
index 6f19a4d..ad412d5 100644
--- a/amulet/sentry.py
+++ b/amulet/sentry.py
@@ -11,6 +11,7 @@
 from . import actions
 from . import waiter
 from . import helpers
+from . import storage
 
 JUJU_VERSION = helpers.JUJU_VERSION
 
@@ -82,6 +83,18 @@ def fromunitdata(cls, unit, unit_data):
         unitsentry.upload_scripts()
         return unitsentry
 
+    def add_storage(self, storage_name, pool=None, count=None, size=None):
+        """Add storage to this unit.
+
+        :param storage_name: Storage name
+        :param pool: Storage pool name
+        :param count: Storage instance count
+        :param size: Storage instance minimum size
+
+        """
+        return storage.add_storage(
+            self.info['unit_name'], storage_name, pool, count, size)
+
     def list_actions(self):
         """Return list of actions defined for this unit.
 
diff --git a/amulet/storage.py b/amulet/storage.py
new file mode 100644
index 0000000..164b771
--- /dev/null
+++ b/amulet/storage.py
@@ -0,0 +1,30 @@
+from .helpers import (
+    juju,
+    JUJU_VERSION,
+)
+
+
+def add_storage(unit, name, pool=None, count=None, size=None):
+    """Add storage to a unit.
+
+    :param unit: Unit on which to add storage, e.g. "wordpress/0"
+    :param name: Storage name
+    :param pool: Storage pool name
+    :param count: Storage instance count
+    :param size: Storage instance minimum size
+
+    """
+    if '/' not in unit:
+        raise ValueError('%s is not a unit' % unit)
+
+    if JUJU_VERSION.major == 1:
+        cmd = ['storage', 'add']
+    else:
+        cmd = ['add-storage']
+
+    constraints = ','.join([
+        str(constraint or '') for constraint in (pool, count, size)])
+    cmd.append(unit)
+    cmd.append('{}={}'.format(name, constraints))
+
+    return juju(cmd)
diff --git a/tests/test_storage.py b/tests/test_storage.py
new file mode 100644
index 0000000..f02bdba
--- /dev/null
+++ b/tests/test_storage.py
@@ -0,0 +1,39 @@
+import unittest
+
+from mock import patch
+
+from amulet import storage
+
+
+class AddStorageTest(unittest.TestCase):
+    def test_bad_unit(self):
+        with self.assertRaises(ValueError):
+            storage.add_storage('foo', 'mystorage')
+
+    @patch('amulet.storage.juju')
+    @patch('amulet.storage.JUJU_VERSION')
+    def test_juju_1(self, version, juju):
+        version.major = 1
+        storage.add_storage(
+            'unit/0',
+            'mystorage',
+            pool='ebs',
+            count=1,
+            size=1024,
+        )
+        juju.assert_called_once_with([
+            'storage', 'add', 'unit/0', 'mystorage=ebs,1,1024'])
+
+    @patch('amulet.storage.juju')
+    @patch('amulet.storage.JUJU_VERSION')
+    def test_juju_2(self, version, juju):
+        version.major = 2
+        storage.add_storage(
+            'unit/0',
+            'mystorage',
+            pool='ebs',
+            count=1,
+            size=1024,
+        )
+        juju.assert_called_once_with([
+            'add-storage', 'unit/0', 'mystorage=ebs,1,1024'])