From 3a3ec9e31d199b936e01530dbc3ea8559d7a82e5 Mon Sep 17 00:00:00 2001
From: Alex Lowe <alex.lowe@canonical.com>
Date: Mon, 8 May 2023 10:49:52 -0400
Subject: [PATCH] bases: Apply stricter typing to base names.

---
 craft_providers/bases/__init__.py | 38 ++++++++++++++++++++-----------
 craft_providers/lxd/remotes.py    |  4 ++--
 tests/unit/bases/test_get_base.py |  2 +-
 tests/unit/lxd/test_remotes.py    |  4 +++-
 4 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/craft_providers/bases/__init__.py b/craft_providers/bases/__init__.py
index b93bd04c..c0fe8709 100644
--- a/craft_providers/bases/__init__.py
+++ b/craft_providers/bases/__init__.py
@@ -19,8 +19,7 @@
 
 # Backward compatible, will be removed in 2.0
 import sys
-from enum import Enum
-from typing import Dict, Tuple, Type
+from typing import Dict, NamedTuple, Tuple, Type, Union
 
 from craft_providers.errors import BaseCompatibilityError, BaseConfigurationError
 
@@ -32,31 +31,44 @@
 
 sys.modules["craft_providers.bases.buildd"] = buildd
 
+BaseAlias = Union[ubuntu.BuilddBaseAlias, centos.CentOSBaseAlias]
+
 __all__ = [
     "ubuntu",
     "centos",
+    "BaseAlias",
+    "BaseName",
     "BuilddBase",
     "BuilddBaseAlias",
     "BaseCompatibilityError",
     "BaseConfigurationError",
 ]
 
-BASE_NAME_TO_BASE_ALIAS: Dict[Tuple[str, str], Enum] = {
-    ("ubuntu", "16.04"): ubuntu.BuilddBaseAlias.XENIAL,
-    ("ubuntu", "18.04"): ubuntu.BuilddBaseAlias.BIONIC,
-    ("ubuntu", "20.04"): ubuntu.BuilddBaseAlias.FOCAL,
-    ("ubuntu", "22.04"): ubuntu.BuilddBaseAlias.JAMMY,
-    ("ubuntu", "22.10"): ubuntu.BuilddBaseAlias.KINETIC,
-    ("ubuntu", "23.04"): ubuntu.BuilddBaseAlias.LUNAR,
-    ("ubuntu", "devel"): ubuntu.BuilddBaseAlias.DEVEL,
-    ("centos", "7"): centos.CentOSBaseAlias.SEVEN,
+
+class BaseName(NamedTuple):
+    """A base image, by distribution and version."""
+
+    name: str
+    version: str
+
+
+BASE_NAME_TO_BASE_ALIAS: Dict[BaseName, BaseAlias] = {
+    BaseName("ubuntu", "16.04"): ubuntu.BuilddBaseAlias.XENIAL,
+    BaseName("ubuntu", "18.04"): ubuntu.BuilddBaseAlias.BIONIC,
+    BaseName("ubuntu", "20.04"): ubuntu.BuilddBaseAlias.FOCAL,
+    BaseName("ubuntu", "22.04"): ubuntu.BuilddBaseAlias.JAMMY,
+    BaseName("ubuntu", "22.10"): ubuntu.BuilddBaseAlias.KINETIC,
+    BaseName("ubuntu", "23.04"): ubuntu.BuilddBaseAlias.LUNAR,
+    BaseName("ubuntu", "devel"): ubuntu.BuilddBaseAlias.DEVEL,
+    BaseName("centos", "7"): centos.CentOSBaseAlias.SEVEN,
 }
 
 
 def get_base_alias(
     base_name: Tuple[str, str],
-) -> Enum:
+) -> BaseAlias:
     """Return a Base alias from a base (name, version) tuple."""
+    base_name = BaseName(*base_name)
     if base_name in BASE_NAME_TO_BASE_ALIAS:
         return BASE_NAME_TO_BASE_ALIAS[base_name]
 
@@ -64,7 +76,7 @@ def get_base_alias(
 
 
 def get_base_from_alias(
-    alias: Enum,
+    alias: BaseAlias,
 ) -> Type[Base]:
     """Return a Base class from a known base alias."""
     if isinstance(alias, ubuntu.BuilddBaseAlias):
diff --git a/craft_providers/lxd/remotes.py b/craft_providers/lxd/remotes.py
index ea834a07..b8e1edee 100644
--- a/craft_providers/lxd/remotes.py
+++ b/craft_providers/lxd/remotes.py
@@ -25,7 +25,7 @@
 from typing import Dict, Union
 
 from craft_providers import Base
-from craft_providers.bases import centos, get_base_alias, ubuntu
+from craft_providers.bases import BaseName, centos, get_base_alias, ubuntu
 
 from .errors import LXDError
 from .lxc import LXC
@@ -180,7 +180,7 @@ def get_remote_image(provider_base: Union[Base, str]) -> RemoteImage:
     """
     # temporary backward compatibility before 2.0
     if isinstance(provider_base, str):
-        alias = get_base_alias(("ubuntu", provider_base))
+        alias = get_base_alias(BaseName("ubuntu", provider_base))
         provider_base = ubuntu.BuilddBase(alias=alias)  # type: ignore
 
     image = _PROVIDER_BASE_TO_LXD_REMOTE_IMAGE.get(provider_base.alias)
diff --git a/tests/unit/bases/test_get_base.py b/tests/unit/bases/test_get_base.py
index 912aa4e2..22be8421 100644
--- a/tests/unit/bases/test_get_base.py
+++ b/tests/unit/bases/test_get_base.py
@@ -30,7 +30,7 @@ def test_get_base_alias_does_not_exist():
         get_base_alias(("ubuntu", "8.04"))
 
     assert exc_info.value == BaseConfigurationError(
-        brief="Base alias not found for ('ubuntu', '8.04')"
+        brief="Base alias not found for BaseName(name='ubuntu', version='8.04')"
     )
 
 
diff --git a/tests/unit/lxd/test_remotes.py b/tests/unit/lxd/test_remotes.py
index ec88788a..00a9bc8d 100644
--- a/tests/unit/lxd/test_remotes.py
+++ b/tests/unit/lxd/test_remotes.py
@@ -210,7 +210,9 @@ def test_get_image_remote_deprecated_error():
     with pytest.raises(BaseConfigurationError) as raised:
         lxd.remotes.get_remote_image("8.04")
 
-    assert "Base alias not found for ('ubuntu', '8.04')" == str(raised.value)
+    assert "Base alias not found for BaseName(name='ubuntu', version='8.04')" == str(
+        raised.value
+    )
 
 
 def test_configure_buildd_image_remote(