Merge pull request canonical#626 from mwhudson/mirror-tweaks

refactor how mirror data is handled

subiquity/controllers/mirror.py

@@ -92,21 +92,21 @@ def start_ui(self):
             self.done(self.answers['mirror'])
         elif 'country-code' in self.answers \
              or 'accept-default' in self.answers:
-            self.done(self.model.mirror)
+            self.done(self.model.get_mirror())
 
     def cancel(self):
         self.app.prev_screen()
 
     def serialize(self):
-        return self.model.mirror
+        return self.model.get_mirror()
 
     def deserialize(self, data):
         super().deserialize(data)
-        self.model.mirror = data
+        self.model.set_mirror(data)
 
     def done(self, mirror):
         log.debug("MirrorController.done next_screen mirror=%s", mirror)
-        if mirror != self.model.mirror:
-            self.model.mirror = mirror
+        if mirror != self.model.get_mirror():
+            self.model.set_mirror(mirror)
         self.configured()
         self.app.next_screen()

subiquity/models/mirror.py

@@ -13,35 +13,59 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import copy
 import logging
-import platform
 from urllib import parse
 
+from curtin.commands.apt_config import (
+    get_arch_mirrorconfig,
+    get_mirror,
+    PRIMARY_ARCHES,
+    )
+from curtin.util import get_architecture
+
 log = logging.getLogger('subiquitycore.models.mirror')
 
-# correct default mirror for most arches
-DEFAULT_MIRROR = 'http://ports.ubuntu.com/ubuntu-ports'
-# apart from the two snowflakes
-if platform.machine() in ['i686', 'x86_64']:
-    DEFAULT_MIRROR = 'http://archive.ubuntu.com/ubuntu'
+
+DEFAULT = {
+    "preserve_sources_list": False,
+    "primary": [
+        {
+            "arches": PRIMARY_ARCHES,
+            "uri": "http://archive.ubuntu.com/ubuntu",
+        },
+        {
+            "arches": ["default"],
+            "uri": "http://ports.ubuntu.com/ubuntu-ports",
+        },
+        ],
+}
 
 
 class MirrorModel(object):
 
     def __init__(self):
-        self.mirror = DEFAULT_MIRROR
+        self.config = copy.deepcopy(DEFAULT)
+        self.architecture = get_architecture()
+        self.default_mirror = self.get_mirror()
 
     def set_country(self, cc):
-        parsed = parse.urlparse(DEFAULT_MIRROR)
+        uri = self.get_mirror()
+        if uri != self.default_mirror:
+            return
+        parsed = parse.urlparse(uri)
         new = parsed._replace(netloc=cc + '.' + parsed.netloc)
-        self.mirror = parse.urlunparse(new)
+        self.set_mirror(parse.urlunparse(new))
+
+    def get_mirror(self):
+        return get_mirror(self.config, "primary", self.architecture)
+
+    def set_mirror(self, mirror):
+        config = get_arch_mirrorconfig(
+            self.config, "primary", self.architecture)
+        config["uri"] = mirror
 
     def render(self):
         return {
-             'apt': {
-                 'primary': [{
-                     'arches': ["default"],
-                     'uri': self.mirror,
-                     }],
-                 }
+             'apt': copy.deepcopy(self.config)
             }

subiquity/models/subiquity.py

@@ -207,10 +207,6 @@ def _machine_id(self):
 
     def render(self, syslog_identifier):
         config = {
-            'apt': {
-                'preserve_sources_list': False,
-                },
-
             'sources': {
                 'ubuntu00': 'cp:///media/filesystem'
                 },

subiquity/models/tests/test_mirror.py

@@ -0,0 +1,44 @@
+# Copyright 2019 Canonical, Ltd.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import unittest
+
+from subiquity.models.mirror import (
+    MirrorModel,
+    )
+
+
+class TestMirrorModel(unittest.TestCase):
+
+    def test_set_country(self):
+        model = MirrorModel()
+        model.set_country("CC")
+        self.assertIn(
+            model.get_mirror(),
+            [
+                "http://CC.archive.ubuntu.com/ubuntu",
+                "http://CC.ports.ubuntu.com/ubuntu-ports",
+            ])
+
+    def test_set_mirror(self):
+        model = MirrorModel()
+        model.set_mirror("http://mymirror.invalid/")
+        self.assertEqual(model.get_mirror(), "http://mymirror.invalid/")
+
+    def test_set_country_after_set_mirror(self):
+        model = MirrorModel()
+        model.set_mirror("http://mymirror.invalid/")
+        model.set_country("CC")
+        self.assertEqual(model.get_mirror(), "http://mymirror.invalid/")

subiquity/models/tests/test_subiquity.py

@@ -129,9 +129,11 @@ def test_has_sources(self):
 
     def test_mirror(self):
         model = SubiquityModel('test')
-        mirror_val = model.mirror.mirror = 'http://my-mirror'
+        mirror_val = 'http://my-mirror'
+        model.mirror.set_mirror(mirror_val)
         config = model.render('ident')
-        val = self.configVal(config, 'apt.primary')
-        self.assertEqual(len(val), 1)
-        val = val[0]
-        self.assertEqual(val['uri'], mirror_val)
+        from curtin.commands.apt_config import get_mirror
+        from curtin.util import get_architecture
+        self.assertEqual(
+            get_mirror(config["apt"], "primary", get_architecture()),
+            mirror_val)

subiquity/ui/views/mirror.py

@@ -50,7 +50,7 @@ def __init__(self, model, controller):
         self.model = model
         self.controller = controller
 
-        self.form = MirrorForm(initial={'url': self.model.mirror})
+        self.form = MirrorForm(initial={'url': self.model.get_mirror()})
 
         connect_signal(self.form, 'submit', self.done)
         connect_signal(self.form, 'cancel', self.cancel)