← Back to branch summary

~le-charmers/charms/trusty/rabbitmq-server/leadership-election

« back to all changes in this revision

Viewing changes to tests/charmhelpers/core/hookenv.py

  • Committer: Liam Young
  • Date: 2015-05-11 08:03:57 UTC
  • mfrom: (83.1.14 rabbitmq-server)
  • Revision ID: liam.young@canonical.com-20150511080357-3ftop9kxb6o0e3mq
Merged trunk in + LE charmhelper sync

Show diffs side-by-side

added added

removed removed

Lines of Context:
tests/charmhelpers/core/hookenv.py
 
1
# Copyright 2014-2015 Canonical Limited.
 
2
#
 
3
# This file is part of charm-helpers.
 
4
#
 
5
# charm-helpers is free software: you can redistribute it and/or modify
 
6
# it under the terms of the GNU Lesser General Public License version 3 as
 
7
# published by the Free Software Foundation.
 
8
#
 
9
# charm-helpers is distributed in the hope that it will be useful,
 
10
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
11
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
12
# GNU Lesser General Public License for more details.
 
13
#
 
14
# You should have received a copy of the GNU Lesser General Public License
 
15
# along with charm-helpers.  If not, see <http://www.gnu.org/licenses/>.
 
16
 
 
17
"Interactions with the Juju environment"
 
18
# Copyright 2013 Canonical Ltd.
 
19
#
 
20
# Authors:
 
21
#  Charm Helpers Developers <juju@lists.ubuntu.com>
 
22
 
 
23
from __future__ import print_function
 
24
from functools import wraps
 
25
import os
 
26
import json
 
27
import yaml
 
28
import subprocess
 
29
import sys
 
30
import errno
 
31
from subprocess import CalledProcessError
 
32
 
 
33
import six
 
34
if not six.PY3:
 
35
    from UserDict import UserDict
 
36
else:
 
37
    from collections import UserDict
 
38
 
 
39
CRITICAL = "CRITICAL"
 
40
ERROR = "ERROR"
 
41
WARNING = "WARNING"
 
42
INFO = "INFO"
 
43
DEBUG = "DEBUG"
 
44
MARKER = object()
 
45
 
 
46
cache = {}
 
47
 
 
48
 
 
49
def cached(func):
 
50
    """Cache return values for multiple executions of func + args
 
51
 
 
52
    For example::
 
53
 
 
54
        @cached
 
55
        def unit_get(attribute):
 
56
            pass
 
57
 
 
58
        unit_get('test')
 
59
 
 
60
    will cache the result of unit_get + 'test' for future calls.
 
61
    """
 
62
    @wraps(func)
 
63
    def wrapper(*args, **kwargs):
 
64
        global cache
 
65
        key = str((func, args, kwargs))
 
66
        try:
 
67
            return cache[key]
 
68
        except KeyError:
 
69
            pass  # Drop out of the exception handler scope.
 
70
        res = func(*args, **kwargs)
 
71
        cache[key] = res
 
72
        return res
 
73
    return wrapper
 
74
 
 
75
 
 
76
def flush(key):
 
77
    """Flushes any entries from function cache where the
 
78
    key is found in the function+args """
 
79
    flush_list = []
 
80
    for item in cache:
 
81
        if key in item:
 
82
            flush_list.append(item)
 
83
    for item in flush_list:
 
84
        del cache[item]
 
85
 
 
86
 
 
87
def log(message, level=None):
 
88
    """Write a message to the juju log"""
 
89
    command = ['juju-log']
 
90
    if level:
 
91
        command += ['-l', level]
 
92
    if not isinstance(message, six.string_types):
 
93
        message = repr(message)
 
94
    command += [message]
 
95
    # Missing juju-log should not cause failures in unit tests
 
96
    # Send log output to stderr
 
97
    try:
 
98
        subprocess.call(command)
 
99
    except OSError as e:
 
100
        if e.errno == errno.ENOENT:
 
101
            if level:
 
102
                message = "{}: {}".format(level, message)
 
103
            message = "juju-log: {}".format(message)
 
104
            print(message, file=sys.stderr)
 
105
        else:
 
106
            raise
 
107
 
 
108
 
 
109
class Serializable(UserDict):
 
110
    """Wrapper, an object that can be serialized to yaml or json"""
 
111
 
 
112
    def __init__(self, obj):
 
113
        # wrap the object
 
114
        UserDict.__init__(self)
 
115
        self.data = obj
 
116
 
 
117
    def __getattr__(self, attr):
 
118
        # See if this object has attribute.
 
119
        if attr in ("json", "yaml", "data"):
 
120
            return self.__dict__[attr]
 
121
        # Check for attribute in wrapped object.
 
122
        got = getattr(self.data, attr, MARKER)
 
123
        if got is not MARKER:
 
124
            return got
 
125
        # Proxy to the wrapped object via dict interface.
 
126
        try:
 
127
            return self.data[attr]
 
128
        except KeyError:
 
129
            raise AttributeError(attr)
 
130
 
 
131
    def __getstate__(self):
 
132
        # Pickle as a standard dictionary.
 
133
        return self.data
 
134
 
 
135
    def __setstate__(self, state):
 
136
        # Unpickle into our wrapper.
 
137
        self.data = state
 
138
 
 
139
    def json(self):
 
140
        """Serialize the object to json"""
 
141
        return json.dumps(self.data)
 
142
 
 
143
    def yaml(self):
 
144
        """Serialize the object to yaml"""
 
145
        return yaml.dump(self.data)
 
146
 
 
147
 
 
148
def execution_environment():
 
149
    """A convenient bundling of the current execution context"""
 
150
    context = {}
 
151
    context['conf'] = config()
 
152
    if relation_id():
 
153
        context['reltype'] = relation_type()
 
154
        context['relid'] = relation_id()
 
155
        context['rel'] = relation_get()
 
156
    context['unit'] = local_unit()
 
157
    context['rels'] = relations()
 
158
    context['env'] = os.environ
 
159
    return context
 
160
 
 
161
 
 
162
def in_relation_hook():
 
163
    """Determine whether we're running in a relation hook"""
 
164
    return 'JUJU_RELATION' in os.environ
 
165
 
 
166
 
 
167
def relation_type():
 
168
    """The scope for the current relation hook"""
 
169
    return os.environ.get('JUJU_RELATION', None)
 
170
 
 
171
 
 
172
def relation_id():
 
173
    """The relation ID for the current relation hook"""
 
174
    return os.environ.get('JUJU_RELATION_ID', None)
 
175
 
 
176
 
 
177
def local_unit():
 
178
    """Local unit ID"""
 
179
    return os.environ['JUJU_UNIT_NAME']
 
180
 
 
181
 
 
182
def remote_unit():
 
183
    """The remote unit for the current relation hook"""
 
184
    return os.environ.get('JUJU_REMOTE_UNIT', None)
 
185
 
 
186
 
 
187
def service_name():
 
188
    """The name service group this unit belongs to"""
 
189
    return local_unit().split('/')[0]
 
190
 
 
191
 
 
192
def hook_name():
 
193
    """The name of the currently executing hook"""
 
194
    return os.path.basename(sys.argv[0])
 
195
 
 
196
 
 
197
class Config(dict):
 
198
    """A dictionary representation of the charm's config.yaml, with some
 
199
    extra features:
 
200
 
 
201
    - See which values in the dictionary have changed since the previous hook.
 
202
    - For values that have changed, see what the previous value was.
 
203
    - Store arbitrary data for use in a later hook.
 
204
 
 
205
    NOTE: Do not instantiate this object directly - instead call
 
206
    ``hookenv.config()``, which will return an instance of :class:`Config`.
 
207
 
 
208
    Example usage::
 
209
 
 
210
        >>> # inside a hook
 
211
        >>> from charmhelpers.core import hookenv
 
212
        >>> config = hookenv.config()
 
213
        >>> config['foo']
 
214
        'bar'
 
215
        >>> # store a new key/value for later use
 
216
        >>> config['mykey'] = 'myval'
 
217
 
 
218
 
 
219
        >>> # user runs `juju set mycharm foo=baz`
 
220
        >>> # now we're inside subsequent config-changed hook
 
221
        >>> config = hookenv.config()
 
222
        >>> config['foo']
 
223
        'baz'
 
224
        >>> # test to see if this val has changed since last hook
 
225
        >>> config.changed('foo')
 
226
        True
 
227
        >>> # what was the previous value?
 
228
        >>> config.previous('foo')
 
229
        'bar'
 
230
        >>> # keys/values that we add are preserved across hooks
 
231
        >>> config['mykey']
 
232
        'myval'
 
233
 
 
234
    """
 
235
    CONFIG_FILE_NAME = '.juju-persistent-config'
 
236
 
 
237
    def __init__(self, *args, **kw):
 
238
        super(Config, self).__init__(*args, **kw)
 
239
        self.implicit_save = True
 
240
        self._prev_dict = None
 
241
        self.path = os.path.join(charm_dir(), Config.CONFIG_FILE_NAME)
 
242
        if os.path.exists(self.path):
 
243
            self.load_previous()
 
244
 
 
245
    def __getitem__(self, key):
 
246
        """For regular dict lookups, check the current juju config first,
 
247
        then the previous (saved) copy. This ensures that user-saved values
 
248
        will be returned by a dict lookup.
 
249
 
 
250
        """
 
251
        try:
 
252
            return dict.__getitem__(self, key)
 
253
        except KeyError:
 
254
            return (self._prev_dict or {})[key]
 
255
 
 
256
    def get(self, key, default=None):
 
257
        try:
 
258
            return self[key]
 
259
        except KeyError:
 
260
            return default
 
261
 
 
262
    def keys(self):
 
263
        prev_keys = []
 
264
        if self._prev_dict is not None:
 
265
            prev_keys = self._prev_dict.keys()
 
266
        return list(set(prev_keys + list(dict.keys(self))))
 
267
 
 
268
    def load_previous(self, path=None):
 
269
        """Load previous copy of config from disk.
 
270
 
 
271
        In normal usage you don't need to call this method directly - it
 
272
        is called automatically at object initialization.
 
273
 
 
274
        :param path:
 
275
 
 
276
            File path from which to load the previous config. If `None`,
 
277
            config is loaded from the default location. If `path` is
 
278
            specified, subsequent `save()` calls will write to the same
 
279
            path.
 
280
 
 
281
        """
 
282
        self.path = path or self.path
 
283
        with open(self.path) as f:
 
284
            self._prev_dict = json.load(f)
 
285
 
 
286
    def changed(self, key):
 
287
        """Return True if the current value for this key is different from
 
288
        the previous value.
 
289
 
 
290
        """
 
291
        if self._prev_dict is None:
 
292
            return True
 
293
        return self.previous(key) != self.get(key)
 
294
 
 
295
    def previous(self, key):
 
296
        """Return previous value for this key, or None if there
 
297
        is no previous value.
 
298
 
 
299
        """
 
300
        if self._prev_dict:
 
301
            return self._prev_dict.get(key)
 
302
        return None
 
303
 
 
304
    def save(self):
 
305
        """Save this config to disk.
 
306
 
 
307
        If the charm is using the :mod:`Services Framework <services.base>`
 
308
        or :meth:'@hook <Hooks.hook>' decorator, this
 
309
        is called automatically at the end of successful hook execution.
 
310
        Otherwise, it should be called directly by user code.
 
311
 
 
312
        To disable automatic saves, set ``implicit_save=False`` on this
 
313
        instance.
 
314
 
 
315
        """
 
316
        if self._prev_dict:
 
317
            for k, v in six.iteritems(self._prev_dict):
 
318
                if k not in self:
 
319
                    self[k] = v
 
320
        with open(self.path, 'w') as f:
 
321
            json.dump(self, f)
 
322
 
 
323
 
 
324
@cached
 
325
def config(scope=None):
 
326
    """Juju charm configuration"""
 
327
    config_cmd_line = ['config-get']
 
328
    if scope is not None:
 
329
        config_cmd_line.append(scope)
 
330
    config_cmd_line.append('--format=json')
 
331
    try:
 
332
        config_data = json.loads(
 
333
            subprocess.check_output(config_cmd_line).decode('UTF-8'))
 
334
        if scope is not None:
 
335
            return config_data
 
336
        return Config(config_data)
 
337
    except ValueError:
 
338
        return None
 
339
 
 
340
 
 
341
@cached
 
342
def relation_get(attribute=None, unit=None, rid=None):
 
343
    """Get relation information"""
 
344
    _args = ['relation-get', '--format=json']
 
345
    if rid:
 
346
        _args.append('-r')
 
347
        _args.append(rid)
 
348
    _args.append(attribute or '-')
 
349
    if unit:
 
350
        _args.append(unit)
 
351
    try:
 
352
        return json.loads(subprocess.check_output(_args).decode('UTF-8'))
 
353
    except ValueError:
 
354
        return None
 
355
    except CalledProcessError as e:
 
356
        if e.returncode == 2:
 
357
            return None
 
358
        raise
 
359
 
 
360
 
 
361
def relation_set(relation_id=None, relation_settings=None, **kwargs):
 
362
    """Set relation information for the current unit"""
 
363
    relation_settings = relation_settings if relation_settings else {}
 
364
    relation_cmd_line = ['relation-set']
 
365
    if relation_id is not None:
 
366
        relation_cmd_line.extend(('-r', relation_id))
 
367
    for k, v in (list(relation_settings.items()) + list(kwargs.items())):
 
368
        if v is None:
 
369
            relation_cmd_line.append('{}='.format(k))
 
370
        else:
 
371
            relation_cmd_line.append('{}={}'.format(k, v))
 
372
    subprocess.check_call(relation_cmd_line)
 
373
    # Flush cache of any relation-gets for local unit
 
374
    flush(local_unit())
 
375
 
 
376
 
 
377
@cached
 
378
def relation_ids(reltype=None):
 
379
    """A list of relation_ids"""
 
380
    reltype = reltype or relation_type()
 
381
    relid_cmd_line = ['relation-ids', '--format=json']
 
382
    if reltype is not None:
 
383
        relid_cmd_line.append(reltype)
 
384
        return json.loads(
 
385
            subprocess.check_output(relid_cmd_line).decode('UTF-8')) or []
 
386
    return []
 
387
 
 
388
 
 
389
@cached
 
390
def related_units(relid=None):
 
391
    """A list of related units"""
 
392
    relid = relid or relation_id()
 
393
    units_cmd_line = ['relation-list', '--format=json']
 
394
    if relid is not None:
 
395
        units_cmd_line.extend(('-r', relid))
 
396
    return json.loads(
 
397
        subprocess.check_output(units_cmd_line).decode('UTF-8')) or []
 
398
 
 
399
 
 
400
@cached
 
401
def relation_for_unit(unit=None, rid=None):
 
402
    """Get the json represenation of a unit's relation"""
 
403
    unit = unit or remote_unit()
 
404
    relation = relation_get(unit=unit, rid=rid)
 
405
    for key in relation:
 
406
        if key.endswith('-list'):
 
407
            relation[key] = relation[key].split()
 
408
    relation['__unit__'] = unit
 
409
    return relation
 
410
 
 
411
 
 
412
@cached
 
413
def relations_for_id(relid=None):
 
414
    """Get relations of a specific relation ID"""
 
415
    relation_data = []
 
416
    relid = relid or relation_ids()
 
417
    for unit in related_units(relid):
 
418
        unit_data = relation_for_unit(unit, relid)
 
419
        unit_data['__relid__'] = relid
 
420
        relation_data.append(unit_data)
 
421
    return relation_data
 
422
 
 
423
 
 
424
@cached
 
425
def relations_of_type(reltype=None):
 
426
    """Get relations of a specific type"""
 
427
    relation_data = []
 
428
    reltype = reltype or relation_type()
 
429
    for relid in relation_ids(reltype):
 
430
        for relation in relations_for_id(relid):
 
431
            relation['__relid__'] = relid
 
432
            relation_data.append(relation)
 
433
    return relation_data
 
434
 
 
435
 
 
436
@cached
 
437
def metadata():
 
438
    """Get the current charm metadata.yaml contents as a python object"""
 
439
    with open(os.path.join(charm_dir(), 'metadata.yaml')) as md:
 
440
        return yaml.safe_load(md)
 
441
 
 
442
 
 
443
@cached
 
444
def relation_types():
 
445
    """Get a list of relation types supported by this charm"""
 
446
    rel_types = []
 
447
    md = metadata()
 
448
    for key in ('provides', 'requires', 'peers'):
 
449
        section = md.get(key)
 
450
        if section:
 
451
            rel_types.extend(section.keys())
 
452
    return rel_types
 
453
 
 
454
 
 
455
@cached
 
456
def charm_name():
 
457
    """Get the name of the current charm as is specified on metadata.yaml"""
 
458
    return metadata().get('name')
 
459
 
 
460
 
 
461
@cached
 
462
def relations():
 
463
    """Get a nested dictionary of relation data for all related units"""
 
464
    rels = {}
 
465
    for reltype in relation_types():
 
466
        relids = {}
 
467
        for relid in relation_ids(reltype):
 
468
            units = {local_unit(): relation_get(unit=local_unit(), rid=relid)}
 
469
            for unit in related_units(relid):
 
470
                reldata = relation_get(unit=unit, rid=relid)
 
471
                units[unit] = reldata
 
472
            relids[relid] = units
 
473
        rels[reltype] = relids
 
474
    return rels
 
475
 
 
476
 
 
477
@cached
 
478
def is_relation_made(relation, keys='private-address'):
 
479
    '''
 
480
    Determine whether a relation is established by checking for
 
481
    presence of key(s).  If a list of keys is provided, they
 
482
    must all be present for the relation to be identified as made
 
483
    '''
 
484
    if isinstance(keys, str):
 
485
        keys = [keys]
 
486
    for r_id in relation_ids(relation):
 
487
        for unit in related_units(r_id):
 
488
            context = {}
 
489
            for k in keys:
 
490
                context[k] = relation_get(k, rid=r_id,
 
491
                                          unit=unit)
 
492
            if None not in context.values():
 
493
                return True
 
494
    return False
 
495
 
 
496
 
 
497
def open_port(port, protocol="TCP"):
 
498
    """Open a service network port"""
 
499
    _args = ['open-port']
 
500
    _args.append('{}/{}'.format(port, protocol))
 
501
    subprocess.check_call(_args)
 
502
 
 
503
 
 
504
def close_port(port, protocol="TCP"):
 
505
    """Close a service network port"""
 
506
    _args = ['close-port']
 
507
    _args.append('{}/{}'.format(port, protocol))
 
508
    subprocess.check_call(_args)
 
509
 
 
510
 
 
511
@cached
 
512
def unit_get(attribute):
 
513
    """Get the unit ID for the remote unit"""
 
514
    _args = ['unit-get', '--format=json', attribute]
 
515
    try:
 
516
        return json.loads(subprocess.check_output(_args).decode('UTF-8'))
 
517
    except ValueError:
 
518
        return None
 
519
 
 
520
 
 
521
def unit_public_ip():
 
522
    """Get this unit's public IP address"""
 
523
    return unit_get('public-address')
 
524
 
 
525
 
 
526
def unit_private_ip():
 
527
    """Get this unit's private IP address"""
 
528
    return unit_get('private-address')
 
529
 
 
530
 
 
531
class UnregisteredHookError(Exception):
 
532
    """Raised when an undefined hook is called"""
 
533
    pass
 
534
 
 
535
 
 
536
class Hooks(object):
 
537
    """A convenient handler for hook functions.
 
538
 
 
539
    Example::
 
540
 
 
541
        hooks = Hooks()
 
542
 
 
543
        # register a hook, taking its name from the function name
 
544
        @hooks.hook()
 
545
        def install():
 
546
            pass  # your code here
 
547
 
 
548
        # register a hook, providing a custom hook name
 
549
        @hooks.hook("config-changed")
 
550
        def config_changed():
 
551
            pass  # your code here
 
552
 
 
553
        if __name__ == "__main__":
 
554
            # execute a hook based on the name the program is called by
 
555
            hooks.execute(sys.argv)
 
556
    """
 
557
 
 
558
    def __init__(self, config_save=True):
 
559
        super(Hooks, self).__init__()
 
560
        self._hooks = {}
 
561
        self._config_save = config_save
 
562
 
 
563
    def register(self, name, function):
 
564
        """Register a hook"""
 
565
        self._hooks[name] = function
 
566
 
 
567
    def execute(self, args):
 
568
        """Execute a registered hook based on args[0]"""
 
569
        hook_name = os.path.basename(args[0])
 
570
        if hook_name in self._hooks:
 
571
            self._hooks[hook_name]()
 
572
            if self._config_save:
 
573
                cfg = config()
 
574
                if cfg.implicit_save:
 
575
                    cfg.save()
 
576
        else:
 
577
            raise UnregisteredHookError(hook_name)
 
578
 
 
579
    def hook(self, *hook_names):
 
580
        """Decorator, registering them as hooks"""
 
581
        def wrapper(decorated):
 
582
            for hook_name in hook_names:
 
583
                self.register(hook_name, decorated)
 
584
            else:
 
585
                self.register(decorated.__name__, decorated)
 
586
                if '_' in decorated.__name__:
 
587
                    self.register(
 
588
                        decorated.__name__.replace('_', '-'), decorated)
 
589
            return decorated
 
590
        return wrapper
 
591
 
 
592
 
 
593
def charm_dir():
 
594
    """Return the root directory of the current charm"""
 
595
    return os.environ.get('CHARM_DIR')
 
596
 
 
597
 
 
598
@cached
 
599
def action_get(key=None):
 
600
    """Gets the value of an action parameter, or all key/value param pairs"""
 
601
    cmd = ['action-get']
 
602
    if key is not None:
 
603
        cmd.append(key)
 
604
    cmd.append('--format=json')
 
605
    action_data = json.loads(subprocess.check_output(cmd).decode('UTF-8'))
 
606
    return action_data
 
607
 
 
608
 
 
609
def action_set(values):
 
610
    """Sets the values to be returned after the action finishes"""
 
611
    cmd = ['action-set']
 
612
    for k, v in list(values.items()):
 
613
        cmd.append('{}={}'.format(k, v))
 
614
    subprocess.check_call(cmd)
 
615
 
 
616
 
 
617
def action_fail(message):
 
618
    """Sets the action status to failed and sets the error message.
 
619
 
 
620
    The results set by action_set are preserved."""
 
621
    subprocess.check_call(['action-fail', message])
 
622
 
 
623
 
 
624
def status_set(workload_state, message):
 
625
    """Set the workload state with a message
 
626
 
 
627
    Use status-set to set the workload state with a message which is visible
 
628
    to the user via juju status. If the status-set command is not found then
 
629
    assume this is juju < 1.23 and juju-log the message unstead.
 
630
 
 
631
    workload_state -- valid juju workload state.
 
632
    message        -- status update message
 
633
    """
 
634
    valid_states = ['maintenance', 'blocked', 'waiting', 'active']
 
635
    if workload_state not in valid_states:
 
636
        raise ValueError(
 
637
            '{!r} is not a valid workload state'.format(workload_state)
 
638
        )
 
639
    cmd = ['status-set', workload_state, message]
 
640
    try:
 
641
        ret = subprocess.call(cmd)
 
642
        if ret == 0:
 
643
            return
 
644
    except OSError as e:
 
645
        if e.errno != errno.ENOENT:
 
646
            raise
 
647
    log_message = 'status-set failed: {} {}'.format(workload_state,
 
648
                                                    message)
 
649
    log(log_message, level='INFO')
 
650
 
 
651
 
 
652
def status_get():
 
653
    """Retrieve the previously set juju workload state
 
654
 
 
655
    If the status-set command is not found then assume this is juju < 1.23 and
 
656
    return 'unknown'
 
657
    """
 
658
    cmd = ['status-get']
 
659
    try:
 
660
        raw_status = subprocess.check_output(cmd, universal_newlines=True)
 
661
        status = raw_status.rstrip()
 
662
        return status
 
663
    except OSError as e:
 
664
        if e.errno == errno.ENOENT:
 
665
            return 'unknown'
 
666
        else:
 
667
            raise