Add iter_units_for_relation_name

Add the iter_units_for_relation_name generator for iterating units
based on relation name.

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

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