Merge lp:~james-page/charm-tools/extra-python-helpers into lp:~charmers/charm-tools/trunk

[James Page]

This MP proposed a number of new helper functions for the python
charm helper.

This wrap a number of juju commands to help with interacting with relations,
retrieving unit and config attributed and interacting with deployed services.

Plus for a bonus a source configurator which uses config attributed 'source'
(and optionally 'key') to configure a charm to use different, non-archive software
sources.

[Mark Mims (mark-mims)]

lgtm

review: Approve

=== modified file 'helpers/python/charmhelpers/__init__.py'
--- helpers/python/charmhelpers/__init__.py	2012-04-11 12:40:22 +0000
+++ helpers/python/charmhelpers/__init__.py	2012-10-02 13:57:18 +0000
@@ -11,6 +11,13 @@
     'log_exit',
     'relation_get',
     'relation_set',
+    'relation_ids',
+    'relation_list',
+    'config_get',
+    'unit_get',
+    'open_port',
+    'close_port',
+    'service_control',
     'unit_info',
     'wait_for_machine',
     'wait_for_page_contents',
@@ -29,6 +36,7 @@
 import time
 import urllib2
 import yaml
+from subprocess import CalledProcessError
 
 
 SLEEP_AMOUNT = 0.1
@@ -48,20 +56,98 @@
 
 
 def get_config():
-    config_get = command('config-get', '--format=json')
-    return json.loads(config_get())
-
-
-def relation_get(*args):
-    cmd = command('relation-get')
-    return cmd(*args).strip()
+    _config_get = command('config-get', '--format=json')
+    return json.loads(_config_get())
+
+def relation_get():
+    cmd = command('relation-get')
+    return cmd().strip()
+
+def relation_get(attribute, unit=None, rid=None):
+    cmd = command('relation-get')
+    _args = []
+    if rid:
+        _args.append('-r')
+        _args.append(rid)
+    _args.append(attribute)
+    if unit:
+        _args.append(unit)
+    return cmd(*_args).strip()
 
 
 def relation_set(**kwargs):
     cmd = command('relation-set')
     args = ['{}={}'.format(k, v) for k, v in kwargs.items()]
-    return cmd(*args)
-
+    cmd(*args)
+
+
+def relation_ids(relation_name):
+    cmd = command('relation-ids')
+    args = [ relation_name ]
+    return cmd(*args).split()
+
+
+def relation_list(rid=None):
+    cmd = command('relation-list')
+    args = []
+    if rid:
+        args.append('-r')
+        args.append(rid)
+    return cmd(*args).split()
+
+
+def config_get(attribute):
+    cmd = command('config-get')
+    args = [ attribute ]
+    return cmd(*args).strip()
+
+
+def unit_get(attribute):
+    cmd = command('unit-get')
+    args = [ attribute ]
+    return cmd(*args).strip()
+
+
+def open_port(port, protocol="TCP"):
+    cmd = command('open-port')
+    args = [ '{}/{}'.format(port, protocol)  ]
+    cmd(*args)
+
+
+def close_port(port, protocol="TCP"):
+    cmd = command('close-port')
+    args = [ '{}/{}'.format(port, protocol)  ]
+    cmd(*args)
+
+START="start"
+RESTART="restart"
+STOP="stop"
+RELOAD="reload"
+
+def service_control(service_name, action):
+    cmd = command('service')
+    args = [ service_name, action ]
+    try:
+        if action == RESTART:
+            try:
+                cmd(*args)
+            except CalledProcessError:
+                service_control(service_name, START)
+        else:
+            cmd(*args)
+    except CalledProcessError:
+        log("Failed to perform {} on service {}".format(action, service_name))
+
+def configure_source(update=False):
+    source = config_get('source')
+    if (source.startswith('ppa:') or
+        source.startswith('cloud:') or
+        source.startswith('http:')):
+        run('add-apt-repository', source)
+    if source.startswith("http:"):
+        run('apt-key', 'import', config_get('key'))
+    if update:
+        run('apt-get', 'update')
 
 def make_charm_config_file(charm_config):
     charm_config_file = tempfile.NamedTemporaryFile()

=== modified file 'helpers/python/charmhelpers/tests/test_charmhelpers.py'
--- helpers/python/charmhelpers/tests/test_charmhelpers.py	2012-04-11 13:15:44 +0000
+++ helpers/python/charmhelpers/tests/test_charmhelpers.py	2012-10-02 13:57:18 +0000
@@ -14,6 +14,8 @@
 sys.path.insert(0, 'helpers/python')
 import charmhelpers
 
+from subprocess import CalledProcessError
+
 
 class CharmHelpersTestCase(TestCase):
     """A basic test case for Python charm helpers."""
@@ -91,6 +93,24 @@
         self._patch_command(lambda: dumps(mock_config))
         self.assertEqual(mock_config, charmhelpers.get_config())
 
+    def test_config_get(self):
+        # config_get is used to retrieve individual configuration elements
+        mock_config = {'key': 'value'}
+
+        # Monkey-patch shelltoolbox.command to avoid having to call out
+        # to config-get.
+        self._patch_command(lambda *args: mock_config[args[0]])
+        self.assertEqual(mock_config['key'], charmhelpers.config_get('key'))
+
+    def test_unit_get(self):
+        # unit_get is used to retrieve individual configuration elements
+        mock_config = {'key': 'value'}
+
+        # Monkey-patch shelltoolbox.command to avoid having to call out
+        # to unit-get.
+        self._patch_command(lambda *args: mock_config[args[0]])
+        self.assertEqual(mock_config['key'], charmhelpers.unit_get('key'))
+
     def test_relation_get(self):
         # relation_get returns the value of a given relation variable,
         # as returned by relation-get $VAR.
@@ -102,6 +122,10 @@
         self.assertEqual('bar', charmhelpers.relation_get('foo'))
         self.assertEqual('eggs', charmhelpers.relation_get('spam'))
 
+        self._patch_command(lambda *args: mock_relation_values[args[2]])
+        self.assertEqual('bar', charmhelpers.relation_get('foo','test','test:1'))
+        self.assertEqual('eggs', charmhelpers.relation_get('spam','test','test:1'))
+
     def test_relation_set(self):
         # relation_set calls out to relation-set and passes key=value
         # pairs to it.
@@ -115,6 +139,86 @@
         self.assertEqual('bar', items_set.get('foo'))
         self.assertEqual('eggs', items_set.get('spam'))
 
+    def test_relation_ids(self):
+        # relation_ids returns a list of relations id for the given
+        # named relation
+        mock_relation_ids = {
+            'test' : 'test:1 test:2'
+            }
+        self._patch_command(lambda *args: mock_relation_ids[args[0]])
+        self.assertEqual(mock_relation_ids['test'].split(),
+                         charmhelpers.relation_ids('test'))
+
+    def test_relation_list(self):
+        # relation_list returns a list of unit names either for the current
+        # context or for the provided relation ID
+        mock_unit_names = {
+            'test:1' : 'test/0 test/1 test/2',
+            'test:2' : 'test/3 test/4 test/5'
+            }
+
+        # Patch command for current context use base - context = test:1
+        self._patch_command(lambda: mock_unit_names['test:1'])
+        self.assertEqual(mock_unit_names['test:1'].split(),
+                         charmhelpers.relation_list())
+        # Patch command for provided relation-id
+        self._patch_command(lambda *args: mock_unit_names[args[1]])
+        self.assertEqual(mock_unit_names['test:2'].split(),
+                         charmhelpers.relation_list(rid='test:2'))
+
+    def test_open_close_port(self):
+        # expose calls open-port with port/protocol parameters
+        ports_set = []
+        def mock_open_port(*args):
+            for arg in args:
+                ports_set.append(arg)
+        def mock_close_port(*args):
+            if args[0] in ports_set:
+                ports_set.remove(args[0])
+        # Monkey patch in the open-port mock
+        self._patch_command(mock_open_port)
+        charmhelpers.open_port(80,"TCP")
+        charmhelpers.open_port(90,"UDP")
+        charmhelpers.open_port(100)
+        self.assertTrue("80/TCP" in ports_set)
+        self.assertTrue("90/UDP" in ports_set)
+        self.assertTrue("100/TCP" in ports_set)
+        # Monkey patch in the close-port mock function
+        self._patch_command(mock_close_port)
+        charmhelpers.close_port(80,"TCP")
+        charmhelpers.close_port(90,"UDP")
+        charmhelpers.close_port(100)
+        # ports_set should now be empty
+        self.assertEquals(len(ports_set), 0)
+
+    def test_service_control(self):
+        # Collect commands that have been run
+        commands_set = {}
+        def mock_service(*args):
+            service = args[0]
+            action = args[1]
+            if service not in commands_set:
+                commands_set[service] = []
+            if (len(commands_set[service]) > 1 and
+                commands_set[service][-1] == 'stop' and
+                action == 'restart'):
+                # Service is stopped - so needs 'start'
+                # action as restart will fail
+                commands_set[service].append(action)
+                raise CalledProcessError(1, repr(args))
+            else:
+                commands_set[service].append(action)
+
+        result = [ 'start', 'stop', 'restart', 'start' ]
+
+        # Monkey patch service command
+        self._patch_command(mock_service)
+        charmhelpers.service_control('myservice','start')
+        charmhelpers.service_control('myservice','stop')
+        charmhelpers.service_control('myservice','restart')
+        self.assertEquals(result, commands_set['myservice'])
+
+
     def test_make_charm_config_file(self):
         # make_charm_config_file() writes the passed configuration to a
         # temporary file as YAML.