From 10a2e6ac227962894631d466fa69ec2655cba58c Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Wed, 15 Feb 2023 14:00:43 +0100
Subject: [PATCH] mirror: allow to combine different filters for the candidates

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 subiquity/models/mirror.py | 57 ++++++++++++++++++++++++++++----------
 1 file changed, 42 insertions(+), 15 deletions(-)

diff --git a/subiquity/models/mirror.py b/subiquity/models/mirror.py
index bcfb2062..765a2fb9 100644
--- a/subiquity/models/mirror.py
+++ b/subiquity/models/mirror.py
@@ -75,7 +75,10 @@ import abc
 import copy
 import contextlib
 import logging
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Union
+from typing import (
+    Any, Callable, Dict, Iterator, List,
+    Optional, Sequence, Set, Union,
+    )
 from urllib import parse
 
 import attr
@@ -130,6 +133,11 @@ class BasePrimaryEntry(abc.ABC):
     def serialize_for_ai(self) -> Any:
         """ Serialize the entry for autoinstall. """
 
+    @abc.abstractmethod
+    def supports_arch(self, arch: str) -> bool:
+        """ Tells whether the mirror claims to support the architecture
+        specified. """
+
 
 @attr.s(auto_attribs=True)
 class PrimaryEntry(BasePrimaryEntry):
@@ -159,8 +167,6 @@ class PrimaryEntry(BasePrimaryEntry):
         return [{"uri": self.uri, "arches": ["default"]}]
 
     def supports_arch(self, arch: str) -> bool:
-        """ Tells whether the mirror claims to support the architecture
-        specified. """
         if self.arches is None:
             return True
         return arch in self.arches
@@ -207,6 +213,11 @@ class LegacyPrimaryEntry(BasePrimaryEntry):
     def serialize_for_ai(self) -> List[Any]:
         return self.config
 
+    def supports_arch(self, arch: str) -> bool:
+        # Curtin will always find a mirror ; albeit with the ["default"]
+        # architectures.
+        return True
+
 
 def countrify_uri(uri: str, cc: str) -> str:
     """ Return a URL where the host is prefixed with a country code. """
@@ -215,6 +226,18 @@ def countrify_uri(uri: str, cc: str) -> str:
     return parse.urlunparse(new)
 
 
+CandidateFilter = Callable[[BasePrimaryEntry], bool]
+
+
+def filter_candidates(candidates: List[BasePrimaryEntry],
+                      *, filters: Sequence[CandidateFilter]) \
+                              -> Iterator[BasePrimaryEntry]:
+    candidates_iter = iter(candidates)
+    for filt in filters:
+        candidates_iter = filter(filt, candidates_iter)
+    return candidates_iter
+
+
 class MirrorModel(object):
 
     def __init__(self):
@@ -312,8 +335,12 @@ class MirrorModel(object):
         # install. It will be placed in etc/apt/sources.list of the target
         # system.
         with contextlib.suppress(StopIteration):
-            candidate = next(filter(lambda c: c.uri is not None,
-                                    self.compatible_primary_candidates()))
+            filters = [
+                lambda c: c.uri is not None,
+                lambda c: c.supports_arch(self.architecture),
+            ]
+            candidate = next(filter_candidates(self.primary_candidates,
+                                               filters=filters))
             return self._get_apt_config_using_candidate(candidate)
         # Our last resort is to include no primary section. Curtin will use
         # its own internal values.
@@ -357,20 +384,20 @@ class MirrorModel(object):
         return next(self.country_mirror_candidates(), None) is not None
 
     def country_mirror_candidates(self) -> Iterator[BasePrimaryEntry]:
-        for candidate in self.primary_candidates:
+        def filt(candidate):
             if self.legacy_primary and candidate.mirror_is_default():
-                yield candidate
+                return True
             elif not self.legacy_primary and candidate.country_mirror:
-                yield candidate
+                return True
+            return False
+
+        return filter_candidates(self.primary_candidates, filters=[filt])
 
     def compatible_primary_candidates(self) -> Iterator[BasePrimaryEntry]:
-        for candidate in self.primary_candidates:
-            if self.legacy_primary:
-                yield candidate
-            elif candidate.arches is None:
-                yield candidate
-            elif self.architecture in candidate.arches:
-                yield candidate
+        def filt(candidate):
+            return candidate.supports_arch(self.architecture)
+
+        return filter_candidates(self.primary_candidates, filters=[filt])
 
     def render(self):
         return {}