mirror: call relevant version of apt_config

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>

subiquity/models/mirror.py

@@ -15,7 +15,7 @@
 
 import copy
 import logging
-from typing import Any, List, Set
+from typing import Any, Dict, Iterator, List, Set
 from urllib import parse
 
 from curtin.commands.apt_config import (
@@ -65,6 +65,11 @@ class MirrorModel(object):
             self.primary_elected,
         ]
 
+        self.iter_primary_candidate: Iterator[PrimarySectionConfig] = \
+            iter(self.primary_candidates)
+        self.primary_staged: PrimarySectionConfig = \
+            next(self.iter_primary_candidate)
+
         self.architecture = get_architecture()
         self.default_mirror = self.get_mirror()
 
@@ -77,12 +82,24 @@ class MirrorModel(object):
             self.primary_elected = self.primary_candidates[0]
         merge_config(self.config, data)
 
-    def get_apt_config(self):
+    def _get_apt_config_common(self) -> Dict[str, Any]:
+        assert "disable_components" not in self.config
+        assert "primary" not in self.config
+
         config = copy.deepcopy(self.config)
-        config["primary"] = copy.deepcopy(self.primary_elected)
         config["disable_components"] = sorted(self.disabled_components)
         return config
 
+    def get_apt_config_staged(self) -> Dict[str, Any]:
+        config = self._get_apt_config_common()
+        config["primary"] = self.primary_staged
+        return config
+
+    def get_apt_config_elected(self) -> Dict[str, Any]:
+        config = self._get_apt_config_common()
+        config["primary"] = self.primary_elected
+        return config
+
     def mirror_is_default(self):
         return self.get_mirror() == self.default_mirror
 
@@ -115,4 +132,6 @@ class MirrorModel(object):
         return {}
 
     def make_autoinstall(self):
-        return self.get_apt_config()
+        config = self._get_apt_config_common()
+        config["primary"] = self.primary_elected
+        return config

subiquity/models/tests/test_mirror.py

@@ -72,20 +72,20 @@ class TestMirrorModel(unittest.TestCase):
         self.assertEqual(self.model.get_mirror(), "http://mymirror.invalid/")
 
     def test_default_disable_components(self):
-        config = self.model.get_apt_config()
+        config = self.model.get_apt_config_staged()
         self.assertEqual([], config['disable_components'])
 
     def test_from_autoinstall(self):
         # autoinstall loads to the config directly
         data = {'disable_components': ['non-free']}
         self.model.load_autoinstall_data(data)
-        config = self.model.get_apt_config()
+        config = self.model.get_apt_config_staged()
         self.assertEqual(['non-free'], config['disable_components'])
 
     def test_disable_add(self):
         expected = ['things', 'stuff']
         self.model.disable_components(expected.copy(), add=True)
-        actual = self.model.get_apt_config()['disable_components']
+        actual = self.model.get_apt_config_staged()['disable_components']
         actual.sort()
         expected.sort()
         self.assertEqual(expected, actual)
@@ -95,7 +95,7 @@ class TestMirrorModel(unittest.TestCase):
         to_remove = ['things', 'stuff']
         expected = ['a', 'b']
         self.model.disable_components(to_remove, add=False)
-        actual = self.model.get_apt_config()['disable_components']
+        actual = self.model.get_apt_config_staged()['disable_components']
         actual.sort()
         expected.sort()
         self.assertEqual(expected, actual)
@@ -103,6 +103,7 @@ class TestMirrorModel(unittest.TestCase):
     def test_make_autoinstall(self):
         primary = [{"arches": "amd64", "uri": "http://mirror"}]
         self.model.disabled_components = set(["non-free"])
+        self.model.primary_candidates = [primary]
         self.model.primary_elected = primary
         cfg = self.model.make_autoinstall()
         self.assertEqual(cfg["disable_components"], ["non-free"])

subiquity/server/apt.py

@@ -112,18 +112,23 @@ class AptConfigurer:
         self.install_tree: Optional[OverlayMountpoint] = None
         self.install_mount = None
 
-    def apt_config(self):
+    def apt_config(self, elected: bool):
         cfg = {}
-        merge_config(cfg, self.app.base_model.mirror.get_apt_config())
+        if elected:
+            merge_config(cfg,
+                         self.app.base_model.mirror.get_apt_config_elected())
+        else:
+            merge_config(cfg,
+                         self.app.base_model.mirror.get_apt_config_staged())
         merge_config(cfg, self.app.base_model.proxy.get_apt_config())
         return {'apt': cfg}
 
-    async def apply_apt_config(self, context):
+    async def apply_apt_config(self, context, elected: bool):
         self.configured_tree = await self.mounter.setup_overlay([self.source])
 
         config_location = os.path.join(
             self.app.root, 'var/log/installer/subiquity-curtin-apt.conf')
-        generate_config_yaml(config_location, self.apt_config())
+        generate_config_yaml(config_location, self.apt_config(elected))
         self.app.note_data_for_apport("CurtinAptConfig", config_location)
 
         await run_curtin_command(

subiquity/server/controllers/mirror.py

@@ -184,11 +184,11 @@ class MirrorController(SubiquityController):
     async def _promote_mirror(self):
         await asyncio.gather(self.source_configured_event.wait(),
                              self.configured_event.wait())
-        await self.apt_configurer.apply_apt_config(self.context)
+        await self.apt_configurer.apply_apt_config(self.context, elected=True)
 
     async def run_mirror_testing(self, output: io.StringIO) -> None:
         await self.source_configured_event.wait()
-        await self.apt_configurer.apply_apt_config(self.context)
+        await self.apt_configurer.apply_apt_config(self.context, elected=False)
         await self.apt_configurer.run_apt_config_check(output)
 
     async def wait_config(self):

subiquity/server/tests/test_apt.py

@@ -67,7 +67,7 @@ class TestAptConfigurer(SubiTestCase):
         self.astart_sym = "subiquity.server.apt.astart_command"
 
     def test_apt_config_noproxy(self):
-        config = self.configurer.apt_config()
+        config = self.configurer.apt_config(elected=False)
         self.assertNotIn("http_proxy", config["apt"])
         self.assertNotIn("https_proxy", config["apt"])
 
@@ -75,7 +75,7 @@ class TestAptConfigurer(SubiTestCase):
         proxy = 'http://apt-cacher-ng:3142'
         self.model.proxy.proxy = proxy
 
-        config = self.configurer.apt_config()
+        config = self.configurer.apt_config(elected=False)
         self.assertEqual(proxy, config["apt"]["http_proxy"])
         self.assertEqual(proxy, config["apt"]["https_proxy"])