Add add_storage() method to UnitSentry

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.
 

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)

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'])