Disable ec2 mirror for non aws instances (#390)

For versions before 20.2, we allowed the use of ec2 mirrors if the datasource availability_zone matches one of the ec2 regions. We are now updating that behavior to allow allow the use of ec2 mirrors on ec2 instances or if the user directly passes an an ec2 mirror url through #cloud-config apt directives.

LP: #1456277

3 files changed, 32 insertions, 8 deletions

diff --git a/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py b/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py
index e5382544..369480be 100644
--- a/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py
+++ b/tests/unittests/test_handler/test_handler_apt_configure_sources_list_v1.py
@@ -101,6 +101,7 @@ class TestAptSourceConfigSourceList(t_help.FilesystemMockingTestCase):
             cfg = {'apt_mirror_search': mirror}
         else:
             cfg = {'apt_mirror': mirror}
+
         mycloud = self._get_cloud(distro)
 
         with mock.patch.object(util, 'write_file') as mockwf:
@@ -108,8 +109,9 @@ class TestAptSourceConfigSourceList(t_help.FilesystemMockingTestCase):
                                    return_value="faketmpl") as mocklf:
                 with mock.patch.object(os.path, 'isfile',
                                        return_value=True) as mockisfile:
-                    with mock.patch.object(templater, 'render_string',
-                                           return_value="fake") as mockrnd:
+                    with mock.patch.object(
+                        templater, 'render_string',
+                            return_value='fake') as mockrnd:
                         with mock.patch.object(util, 'rename'):
                             cc_apt_configure.handle("test", cfg, mycloud,
                                                     LOG, None)
diff --git a/tests/unittests/test_handler/test_handler_apt_source_v1.py b/tests/unittests/test_handler/test_handler_apt_source_v1.py
index f2349157..367971cb 100644
--- a/tests/unittests/test_handler/test_handler_apt_source_v1.py
+++ b/tests/unittests/test_handler/test_handler_apt_source_v1.py
@@ -43,10 +43,17 @@ class FakeDistro(object):
         return
 
 
+class FakeDatasource:
+    """Fake Datasource helper object"""
+    def __init__(self):
+        self.region = 'region'
+
+
 class FakeCloud(object):
     """Fake Cloud helper object"""
     def __init__(self):
         self.distro = FakeDistro()
+        self.datasource = FakeDatasource()
 
 
 class TestAptSourceConfig(TestCase):
diff --git a/tests/unittests/test_handler/test_handler_apt_source_v3.py b/tests/unittests/test_handler/test_handler_apt_source_v3.py
index 220100e2..ac847238 100644
--- a/tests/unittests/test_handler/test_handler_apt_source_v3.py
+++ b/tests/unittests/test_handler/test_handler_apt_source_v3.py
@@ -49,6 +49,18 @@ MOCK_LSB_RELEASE_DATA = {
     'release': '18.04', 'codename': 'bionic'}
 
 
+class FakeDatasource:
+    """Fake Datasource helper object"""
+    def __init__(self):
+        self.region = 'region'
+
+
+class FakeCloud:
+    """Fake Cloud helper object"""
+    def __init__(self):
+        self.datasource = FakeDatasource()
+
+
 class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
     """TestAptSourceConfig
     Main Class to test apt configs
@@ -471,7 +483,7 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
         fromfn = ("%s/%s_%s" % (pre, archive, post))
         tofn = ("%s/test.ubuntu.com_%s" % (pre, post))
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, arch)
+        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, FakeCloud(), arch)
 
         self.assertEqual(mirrors['MIRROR'],
                          "http://test.ubuntu.com/%s/" % component)
@@ -559,7 +571,8 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
                "security": [{'arches': ["default"],
                              "uri": smir}]}
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, 'amd64')
+        mirrors = cc_apt_configure.find_apt_mirror_info(
+            cfg, FakeCloud(), 'amd64')
 
         self.assertEqual(mirrors['MIRROR'],
                          pmir)
@@ -594,7 +607,7 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
                "security": [{'arches': ["default"], "uri": "nothis-security"},
                             {'arches': [arch], "uri": smir}]}
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, arch)
+        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, FakeCloud(), arch)
 
         self.assertEqual(mirrors['PRIMARY'], pmir)
         self.assertEqual(mirrors['MIRROR'], pmir)
@@ -613,7 +626,8 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
                             {'arches': ["default"],
                              "uri": smir}]}
 
-        mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, 'amd64')
+        mirrors = cc_apt_configure.find_apt_mirror_info(
+            cfg, FakeCloud(), 'amd64')
 
         self.assertEqual(mirrors['MIRROR'],
                          pmir)
@@ -673,7 +687,7 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
 
         with mock.patch.object(cc_apt_configure.util, 'search_for_mirror',
                                side_effect=[pmir, smir]) as mocksearch:
-            mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None,
+            mirrors = cc_apt_configure.find_apt_mirror_info(cfg, FakeCloud(),
                                                             'amd64')
 
         calls = [call(["pfailme", pmir]),
@@ -712,7 +726,8 @@ class TestAptSourceConfig(t_help.FilesystemMockingTestCase):
         # should not be called, since primary is specified
         with mock.patch.object(cc_apt_configure.util,
                                'search_for_mirror') as mockse:
-            mirrors = cc_apt_configure.find_apt_mirror_info(cfg, None, arch)
+            mirrors = cc_apt_configure.find_apt_mirror_info(
+                cfg, FakeCloud(), arch)
         mockse.assert_not_called()
 
         self.assertEqual(mirrors['MIRROR'],