From f8f6ca1fc2feb012335f5b276fec43fbd33c595c Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Sun, 29 Jan 2023 13:19:05 +0100
Subject: [PATCH] mirror: add PrimarySection class with tests

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 subiquity/models/mirror.py            | 48 ++++++++++++++++++++++-----
 subiquity/models/tests/test_mirror.py | 27 +++++++++++++++
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/subiquity/models/mirror.py b/subiquity/models/mirror.py
index 93af0695..f17098ad 100644
--- a/subiquity/models/mirror.py
+++ b/subiquity/models/mirror.py
@@ -33,13 +33,50 @@ except ImportError:
 log = logging.getLogger('subiquity.models.mirror')
 
 
+DEFAULT_PRIMARY_SECTION = [
+    {
+        "arches": PRIMARY_ARCHES,
+        "uri": "http://archive.ubuntu.com/ubuntu",
+    }, {
+        "arches": ["default"],
+        "uri": "http://ports.ubuntu.com/ubuntu-ports",
+    },
+]
+
+
 DEFAULT = {
     "preserve_sources_list": False,
 }
 
+
 PrimarySectionConfig = List[Any]
 
 
+class PrimarySection:
+    """ Helper to manage a primary autoinstall section. """
+    def __init__(self, config: List[Any], *, parent: "MirrorModel") -> None:
+        self.parent = parent
+        self.config = config
+
+    def get_mirror(self) -> str:
+        config = copy.deepcopy(self.parent.config)
+        config["primary"] = self.config
+        return get_mirror(config, "primary", self.parent.architecture)
+
+    def set_mirror(self, uri: str) -> None:
+        config = get_arch_mirrorconfig(
+                {"primary": self.config},
+                "primary", self.parent.architecture)
+        config["uri"] = uri
+
+    def mirror_is_default(self) -> bool:
+        return self.get_mirror() == self.parent.default_mirror
+
+    @classmethod
+    def new_from_default(cls, parent: "MirrorModel") -> "PrimarySection":
+        return cls(copy.deepcopy(DEFAULT_PRIMARY_SECTION), parent=parent)
+
+
 def countrify_uri(uri: str, cc: str) -> str:
     """ Return a URL where the host is prefixed with a country code. """
     parsed = parse.urlparse(uri)
@@ -52,15 +89,8 @@ class MirrorModel(object):
     def __init__(self):
         self.config = copy.deepcopy(DEFAULT)
         self.disabled_components: Set[str] = set()
-        self.primary_elected: PrimarySectionConfig = [
-            {
-                "arches": PRIMARY_ARCHES,
-                "uri": "http://archive.ubuntu.com/ubuntu",
-            }, {
-                "arches": ["default"],
-                "uri": "http://ports.ubuntu.com/ubuntu-ports",
-            },
-        ]
+        self.primary_elected: PrimarySectionConfig = \
+            copy.deepcopy(DEFAULT_PRIMARY_SECTION)
         self.primary_candidates: List[PrimarySectionConfig] = [
             self.primary_elected,
         ]
diff --git a/subiquity/models/tests/test_mirror.py b/subiquity/models/tests/test_mirror.py
index 09d4f257..82b5ca68 100644
--- a/subiquity/models/tests/test_mirror.py
+++ b/subiquity/models/tests/test_mirror.py
@@ -17,7 +17,9 @@ import unittest
 
 from subiquity.models.mirror import (
     countrify_uri,
+    DEFAULT_PRIMARY_SECTION,
     MirrorModel,
+    PrimarySection,
     )
 
 
@@ -49,6 +51,31 @@ class TestCountrifyUrl(unittest.TestCase):
                 "http://us.ports.ubuntu.com/ubuntu-ports")
 
 
+class TestPrimarySection(unittest.TestCase):
+    def setUp(self):
+        self.model = MirrorModel()
+
+    def test_initializer(self):
+        primary = PrimarySection([], parent=self.model)
+        self.assertEqual(primary.config, [])
+        self.assertEqual(primary.parent, self.model)
+
+    def test_new_from_default(self):
+        primary = PrimarySection.new_from_default(parent=self.model)
+        self.assertEqual(primary.config, DEFAULT_PRIMARY_SECTION)
+
+    def test_get_mirror(self):
+        self.model.architecture = "amd64"
+        primary = PrimarySection([{"uri": "http://myurl", "arches": "amd64"}],
+                                 parent=self.model)
+        self.assertEqual(primary.get_mirror(), "http://myurl")
+
+    def test_set_mirror(self):
+        primary = PrimarySection.new_from_default(parent=self.model)
+        primary.set_mirror("http://mymirror.invalid/")
+        self.assertEqual(primary.get_mirror(), "http://mymirror.invalid/")
+
+
 class TestMirrorModel(unittest.TestCase):
     def setUp(self):
         self.model = MirrorModel()