diff --git a/charmhelpers/core/hookenv.py b/charmhelpers/core/hookenv.py index bac6d80c8..5a88f798e 100644 --- a/charmhelpers/core/hookenv.py +++ b/charmhelpers/core/hookenv.py @@ -22,6 +22,7 @@ import copy from distutils.version import LooseVersion from functools import wraps +from collections import namedtuple import glob import os import json @@ -1135,28 +1136,6 @@ def network_get(endpoint, relation_id=None): return yaml.safe_load(response) -def ingress_address(rid=None, unit=None): - """ - Retrieve the ingress-address from a relation when available. Otherwise, - return the private-address. This function is to be used on the consuming - side of the relation. - - Usage: - addresses = [] - for rid in relation_ids(relation_name): - for unit in related_units(rid): - addresses.append(ingress_address(rid=rid, unit=unit)) - - :param rid: string relation id - :param unit: string unit name - :side effect: calls relation_get - :return: string IP address - """ - settings = relation_get(rid=rid, unit=unit) - return (settings.get('ingress-address') or - settings.get('private-address')) - - def add_metric(*args, **kwargs): """Add metric values. Values may be expressed with keyword arguments. For metric names containing dashes, these may be expressed as one or more @@ -1186,3 +1165,42 @@ def meter_info(): """Get the meter status information, if running in the meter-status-changed hook.""" return os.environ.get('JUJU_METER_INFO') + + +def iter_units_for_relation_name(relation_name): + """Iterate through all units in a relation + + Generator that iterates through all the units in a relation and yields + a named tuple with rid and unit field names. + + Usage: + data = [(u.rid, u.unit) + for u in iter_units_for_relation_name(relation_name)] + + :param relation_name: string relation name + :yield: Named Tuple with rid and unit field names + """ + RelatedUnit = namedtuple('RelatedUnit', 'rid, unit') + for rid in relation_ids(relation_name): + for unit in related_units(rid): + yield RelatedUnit(rid, unit) + + +def ingress_address(rid=None, unit=None): + """ + Retrieve the ingress-address from a relation when available. Otherwise, + return the private-address. This function is to be used on the consuming + side of the relation. + + Usage: + addresses = [ingress_address(rid=u.rid, unit=u.unit) + for u in iter_units_for_relation_name(relation_name)] + + :param rid: string relation id + :param unit: string unit name + :side effect: calls relation_get + :return: string IP address + """ + settings = relation_get(rid=rid, unit=unit) + return (settings.get('ingress-address') or + settings.get('private-address')) diff --git a/tests/core/test_hookenv.py b/tests/core/test_hookenv.py index ae361be75..d895ca701 100644 --- a/tests/core/test_hookenv.py +++ b/tests/core/test_hookenv.py @@ -1672,27 +1672,6 @@ def test_network_get_primary(self, check_output): ['network-get', '--primary-address', 'mybinding']) self.assertEqual(ip, '192.168.22.1') - @patch.object(hookenv, 'relation_get') - def test_ingress_address(self, relation_get): - """Ensure ingress_address returns the ingress-address when available - and returns the private-address when not. - """ - _with_ingress = {'egress-subnets': '10.5.0.23/32', - 'ingress-address': '10.5.0.23', - 'private-address': '172.16.5.10'} - - _without_ingress = {'private-address': '172.16.5.10'} - - # Return the ingress-address - relation_get.return_value = _with_ingress - self.assertEqual(hookenv.ingress_address(rid='test:1', unit='unit/1'), - '10.5.0.23') - relation_get.assert_called_with(rid='test:1', unit='unit/1') - # Return the private-address - relation_get.return_value = _without_ingress - self.assertEqual(hookenv.ingress_address(rid='test:1'), - '172.16.5.10') - @patch('subprocess.check_output') def test_network_get_primary_unsupported(self, check_output): """Ensure that NotImplementedError is thrown when run on Juju < 2.0""" @@ -1769,3 +1748,37 @@ def test_meter_status(self, os_): } self.assertEqual(hookenv.meter_status(), 'GREEN') self.assertEqual(hookenv.meter_info(), 'all good') + + @patch.object(hookenv, 'related_units') + @patch.object(hookenv, 'relation_ids') + def test_iter_units_for_relation_name(self, relation_ids, related_units): + relation_ids.return_value = ['rel:1'] + related_units.return_value = ['unit/0', 'unit/1', 'unit/2'] + expected = [('rel:1', 'unit/0'), + ('rel:1', 'unit/1'), + ('rel:1', 'unit/2')] + related_units_data = [ + (u.rid, u.unit) + for u in hookenv.iter_units_for_relation_name('rel')] + self.assertEqual(expected, related_units_data) + + @patch.object(hookenv, 'relation_get') + def test_ingress_address(self, relation_get): + """Ensure ingress_address returns the ingress-address when available + and returns the private-address when not. + """ + _with_ingress = {'egress-subnets': '10.5.0.23/32', + 'ingress-address': '10.5.0.23', + 'private-address': '172.16.5.10'} + + _without_ingress = {'private-address': '172.16.5.10'} + + # Return the ingress-address + relation_get.return_value = _with_ingress + self.assertEqual(hookenv.ingress_address(rid='test:1', unit='unit/1'), + '10.5.0.23') + relation_get.assert_called_with(rid='test:1', unit='unit/1') + # Return the private-address + relation_get.return_value = _without_ingress + self.assertEqual(hookenv.ingress_address(rid='test:1'), + '172.16.5.10')