Make full data source available to code that handles mirror selection.

3 files changed, 23 insertions, 17 deletions

diff --git a/cloudinit/distros/__init__.py b/cloudinit/distros/__init__.py
index 8a947867..47b76c68 100644
--- a/cloudinit/distros/__init__.py
+++ b/cloudinit/distros/__init__.py
@@ -117,12 +117,11 @@ class Distro(object):
             arch = self.get_primary_arch()
         return _get_arch_package_mirror_info(mirror_info, arch)
 
-    def get_package_mirror_info(self, arch=None,
-                                availability_zone=None):
+    def get_package_mirror_info(self, arch=None, data_source=None):
         # This resolves the package_mirrors config option
         # down to a single dict of {mirror_name: mirror_url}
         arch_info = self._get_arch_package_mirror_info(arch)
-        return _get_package_mirror_info(availability_zone=availability_zone,
+        return _get_package_mirror_info(data_source=data_source,
                                         mirror_info=arch_info)
 
     def apply_network(self, settings, bring_up=True):
@@ -556,7 +555,7 @@ class Distro(object):
                 LOG.info("Added user '%s' to group '%s'" % (member, name))
 
 
-def _get_package_mirror_info(mirror_info, availability_zone=None,
+def _get_package_mirror_info(mirror_info, data_source=None,
                              mirror_filter=util.search_for_mirror):
     # given a arch specific 'mirror_info' entry (from package_mirrors)
     # search through the 'search' entries, and fallback appropriately
@@ -572,11 +571,11 @@ def _get_package_mirror_info(mirror_info, availability_zone=None,
     ec2_az_re = ("^[a-z][a-z]-(%s)-[1-9][0-9]*[a-z]$" % directions_re)
 
     subst = {}
-    if availability_zone:
-        subst['availability_zone'] = availability_zone
+    if data_source and data_source.availability_zone:
+        subst['availability_zone'] = data_source.availability_zone
 
-    if availability_zone and re.match(ec2_az_re, availability_zone):
-        subst['ec2_region'] = "%s" % availability_zone[0:-1]
+        if re.match(ec2_az_re, data_source.availability_zone):
+            subst['ec2_region'] = "%s" % data_source.availability_zone[0:-1]
 
     results = {}
     for (name, mirror) in mirror_info.get('failsafe', {}).items():
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index 39eab51b..1a036638 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -210,8 +210,7 @@ class DataSource(object):
             return hostname
 
     def get_package_mirror_info(self):
-        return self.distro.get_package_mirror_info(
-            availability_zone=self.availability_zone)
+        return self.distro.get_package_mirror_info(data_source=self)
 
 
 def normalize_pubkey_data(pubkey_data):
diff --git a/tests/unittests/test_distros/test_generic.py b/tests/unittests/test_distros/test_generic.py
index 8e3bd78a..6ed1704c 100644
--- a/tests/unittests/test_distros/test_generic.py
+++ b/tests/unittests/test_distros/test_generic.py
@@ -7,6 +7,11 @@ import os
 import shutil
 import tempfile
 
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
 unknown_arch_info = {
     'arches': ['default'],
     'failsafe': {'primary': 'http://fs-primary-default',
@@ -144,33 +149,35 @@ class TestGenericDistro(helpers.FilesystemMockingTestCase):
 
     def test_get_package_mirror_info_az_ec2(self):
         arch_mirrors = gapmi(package_mirrors, arch="amd64")
+        data_source_mock = mock.Mock(availability_zone="us-east-1a")
 
-        results = gpmi(arch_mirrors, availability_zone="us-east-1a",
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_first)
         self.assertEqual(results,
                          {'primary': 'http://us-east-1.ec2/',
                           'security': 'http://security-mirror1-intel'})
 
-        results = gpmi(arch_mirrors, availability_zone="us-east-1a",
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_second)
         self.assertEqual(results,
                          {'primary': 'http://us-east-1a.clouds/',
                           'security': 'http://security-mirror2-intel'})
 
-        results = gpmi(arch_mirrors, availability_zone="us-east-1a",
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_none)
         self.assertEqual(results, package_mirrors[0]['failsafe'])
 
     def test_get_package_mirror_info_az_non_ec2(self):
         arch_mirrors = gapmi(package_mirrors, arch="amd64")
+        data_source_mock = mock.Mock(availability_zone="nova.cloudvendor")
 
-        results = gpmi(arch_mirrors, availability_zone="nova.cloudvendor",
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_first)
         self.assertEqual(results,
                          {'primary': 'http://nova.cloudvendor.clouds/',
                           'security': 'http://security-mirror1-intel'})
 
-        results = gpmi(arch_mirrors, availability_zone="nova.cloudvendor",
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_last)
         self.assertEqual(results,
                          {'primary': 'http://nova.cloudvendor.clouds/',
@@ -178,17 +185,18 @@ class TestGenericDistro(helpers.FilesystemMockingTestCase):
 
     def test_get_package_mirror_info_none(self):
         arch_mirrors = gapmi(package_mirrors, arch="amd64")
+        data_source_mock = mock.Mock(availability_zone=None)
 
         # because both search entries here replacement based on
         # availability-zone, the filter will be called with an empty list and
         # failsafe should be taken.
-        results = gpmi(arch_mirrors, availability_zone=None,
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_first)
         self.assertEqual(results,
                          {'primary': 'http://fs-primary-intel',
                           'security': 'http://security-mirror1-intel'})
 
-        results = gpmi(arch_mirrors, availability_zone=None,
+        results = gpmi(arch_mirrors, data_source=data_source_mock,
                        mirror_filter=self.return_last)
         self.assertEqual(results,
                          {'primary': 'http://fs-primary-intel',