Replace snap channel config for hard coded revision config

charms/worker/charmcraft.yaml

@@ -42,10 +42,6 @@ bases:
       architectures: [amd64]
 config:
   options:
-    channel:
-      default: edge
-      type: string
-      description: Snap channel of the k8s snap
     labels:
       default: ""
       type: string
@@ -59,15 +55,17 @@ parts:
   charm:
     build-packages: [git]
     charm-entrypoint: k8s/src/charm.py
-  lib:
-    # move the ./k8s/lib path to ./lib since 
-    # charmcraft assumes it to be there once the charm runs
+  promote:
+    # move paths out of ./k8s to ./ since 
+    # charmcraft assumes ./lib to be there
+    # charmcraft assumes ./templates to be there
     after: [charm]
     plugin: nil
     source: ./
     override-prime: |
-      rm -rf $CRAFT_PRIME/lib
+      rm -rf $CRAFT_PRIME/lib $CRAFT_PRIME/templates
       mv $CRAFT_PRIME/k8s/lib $CRAFT_PRIME/lib
+      mv $CRAFT_PRIME/k8s/templates $CRAFT_PRIME/templates
 
 provides:
   cos-agent:

charms/worker/k8s/charmcraft.yaml

@@ -51,10 +51,6 @@ bases:
       architectures: [amd64]
 config:
   options:
-    channel:
-      default: edge
-      type: string
-      description: Snap channel of the k8s snap
     datastore:
       default: dqlite
       type: string

charms/worker/k8s/src/charm.py

@@ -28,6 +28,7 @@
 from urllib.parse import urlparse
 
 import charms.contextual_status as status
+import charms.operator_libs_linux.v2.snap as snap_lib
 import ops
 import yaml
 from charms.contextual_status import WaitingStatus, on_error
@@ -50,10 +51,9 @@
 )
 from charms.kubernetes_libs.v0.etcd import EtcdReactiveRequires
 from charms.node_base import LabelMaker
-from charms.operator_libs_linux.v2.snap import SnapError, SnapState
-from charms.operator_libs_linux.v2.snap import ensure as snap_ensure
 from charms.reconciler import Reconciler
 from cos_integration import COSIntegration
+from snap import management as snap_management
 from token_distributor import ClusterTokenType, TokenCollector, TokenDistributor, TokenStrategy
 
 # Log messages can be retrieved using juju debug-log
@@ -198,12 +198,11 @@ def get_cloud_name(self) -> str:
         """
         return self.xcp.name or ""
 
-    @on_error(ops.BlockedStatus("Failed to install k8s snap."), SnapError)
-    def _install_k8s_snap(self):
-        """Install the k8s snap package."""
+    @on_error(ops.BlockedStatus("Failed to install snaps."), snap_lib.SnapError)
+    def _install_snaps(self):
+        """Install snap packages."""
         status.add(ops.MaintenanceStatus("Ensuring snap installation"))
-        log.info("Ensuring k8s snap version")
-        snap_ensure("k8s", SnapState.Latest.value, self.config["channel"])
+        snap_management()
 
     @on_error(WaitingStatus("Waiting to apply snap requirements"), subprocess.CalledProcessError)
     def _apply_snap_requirements(self):
@@ -527,7 +526,7 @@ def _reconcile(self, event):
             return
 
         self._apply_proxy_environment()
-        self._install_k8s_snap()
+        self._install_snaps()
         self._apply_snap_requirements()
         self._check_k8sd_ready()
         if self.lead_control_plane:

charms/worker/k8s/src/snap.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+# Learn more at: https://juju.is/docs/sdk
+
+"""Snap Installation Module."""
+
+
+import logging
+import subprocess
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+import charms.operator_libs_linux.v2.snap as snap_lib
+import yaml
+from pydantic import BaseModel, Field, ValidationError, parse_obj_as
+from typing_extensions import Annotated
+
+# Log messages can be retrieved using juju debug-log
+log = logging.getLogger(__name__)
+
+
+class SnapFileArgument(BaseModel):
+    """Structure to install a snap by file.
+
+    Attributes:
+        install_type (str): literal string defining this type
+        name (str): The name of the snap after installed
+        filename (Path): Path to the snap to locally install
+        classic (bool): If it should be installed as a classic snap
+        dangerous (bool): If it should be installed as a dangerouse snap
+        devmode (bool): If it should be installed as with dev mode enabled
+    """
+
+    install_type: Literal["file"] = Field("file", alias="install-type", exclude=True)
+    name: str = Field(exclude=True)
+    filename: Optional[Path] = None
+    classic: Optional[bool] = None
+    devmode: Optional[bool] = None
+    dangerous: Optional[bool] = None
+
+
+class SnapStoreArgument(BaseModel):
+    """Structure to install a snap by snapstore.
+
+    Attributes:
+        install_type (str): literal string defining this type
+        name (str): The type of the request.
+        state (SnapState): a `SnapState` to reconcile to.
+        classic (bool): If it should be installed as a classic snap
+        devmode (bool): If it should be installed as with dev mode enabled
+        channel (bool): the channel to install from
+        cohort (str): the key of a cohort that this snap belongs to
+        revision (int): the revision of the snap to install
+    """
+
+    install_type: Literal["store"] = Field("store", alias="install-type", exclude=True)
+    name: str = Field(exclude=True)
+    classic: Optional[bool] = None
+    devmode: Optional[bool] = None
+    state: Optional[snap_lib.SnapState] = Field(snap_lib.SnapState.Present)
+    channel: Optional[str] = None
+    cohort: Optional[str] = None
+    revision: Optional[int] = None
+
+
+SnapArgument = Annotated[
+    Union[SnapFileArgument, SnapStoreArgument], Field(discriminator="install_type")
+]
+
+
+def _parse_management_arguments() -> List[SnapArgument]:
+    """Parse snap management arguments.
+
+    Raises:
+        SnapError: when the management issue cannot be resolved
+
+    Returns:
+        Parsed arguments list for the specific host architecture
+    """
+    revision = Path("templates/snap_installation.yaml")
+    if not revision.exists():
+        raise snap_lib.SnapError(f"Failed to find file={revision}")
+    try:
+        with revision.open() as f:
+            body = yaml.safe_load(f)
+    except yaml.YAMLError as e:
+        log.error("Failed to load file=%s, %s", revision, e)
+        raise snap_lib.SnapError(f"Failed to load file={revision}")
+    dpkg_arch = ["dpkg", "--print-architecture"]
+    arch = subprocess.check_output(dpkg_arch).decode("UTF-8").strip()
+
+    if not (isinstance(body, dict) and (arch_spec := body.get(arch))):
+        log.warning("Failed to find revision for arch=%s", arch)
+        raise snap_lib.SnapError(f"Failed to find revision for arch={arch}")
+
+    try:
+        args: List[SnapArgument] = [parse_obj_as(SnapArgument, arg) for arg in arch_spec]  # type: ignore[arg-type]
+    except ValidationError as e:
+        log.warning("Failed to validate args=%s (%s)", arch_spec, e)
+        raise snap_lib.SnapError("Failed to validate snap args")
+
+    return args
+
+
+def management():
+    """Manage snap installations on this machine."""
+    cache = snap_lib.SnapCache()
+    for args in _parse_management_arguments():
+        which = cache[args.name]
+        if isinstance(args, SnapFileArgument) and which.revision != "x1":
+            snap_lib.install_local(**args.dict(exclude_none=True))
+        elif isinstance(args, SnapStoreArgument):
+            log.info("Ensuring %s snap version", args.name)
+            which.ensure(**args.dict(exclude_none=True))

charms/worker/k8s/templates/snap_installation.yaml

@@ -0,0 +1,6 @@
+amd64:
+- name: k8s
+  channel: edge
+arm64:
+- name: k8s
+  channel: edge

charms/worker/k8s/tests/unit/test_base.py

@@ -42,7 +42,7 @@ def mock_reconciler_handlers(harness):
     """
     handler_names = {
         "_evaluate_removal",
-        "_install_k8s_snap",
+        "_install_snaps",
         "_apply_snap_requirements",
         "_check_k8sd_ready",
         "_join_cluster",

charms/worker/k8s/tests/unit/test_snap.py

@@ -0,0 +1,122 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+# Learn more about testing at: https://juju.is/docs/sdk/testing
+
+# pylint: disable=duplicate-code,missing-function-docstring
+"""Unit tests snap module."""
+
+import io
+import unittest.mock as mock
+from pathlib import Path
+
+import pytest
+import snap
+
+
+@mock.patch("pathlib.Path.exists", mock.Mock(return_value=False))
+def test_parse_no_file():
+    """Test no file exists."""
+    with pytest.raises(snap.snap_lib.SnapError):
+        snap._parse_management_arguments()
+
+
+@mock.patch("pathlib.Path.exists", mock.Mock(return_value=True))
+@mock.patch("pathlib.Path.open")
+def test_parse_invalid_file(mock_open):
+    """Test file is invalid."""
+    mock_open().__enter__.return_value = io.StringIO("example: =")
+    with pytest.raises(snap.snap_lib.SnapError):
+        snap._parse_management_arguments()
+
+
+@mock.patch("pathlib.Path.exists", mock.Mock(return_value=True))
+@mock.patch("pathlib.Path.open")
+@mock.patch("subprocess.check_output")
+def test_parse_invalid_arch(mock_checkoutput, mock_open):
+    """Test file has invalid arch."""
+    mock_open().__enter__.return_value = io.StringIO("{}")
+    mock_checkoutput().decode.return_value = "amd64"
+    with pytest.raises(snap.snap_lib.SnapError):
+        snap._parse_management_arguments()
+
+
+@mock.patch("pathlib.Path.exists", mock.Mock(return_value=True))
+@mock.patch("pathlib.Path.open")
+@mock.patch("subprocess.check_output")
+def test_parse_validation_error(mock_checkoutput, mock_open):
+    """Test file cannot be parsed."""
+    mock_open().__enter__.return_value = io.StringIO("amd64:\n- {}")
+    mock_checkoutput().decode.return_value = "amd64"
+    with pytest.raises(snap.snap_lib.SnapError):
+        snap._parse_management_arguments()
+
+
+@mock.patch("pathlib.Path.exists", mock.Mock(return_value=True))
+@mock.patch("pathlib.Path.open")
+@mock.patch("subprocess.check_output")
+def test_parse_valid_store(mock_checkoutput, mock_open):
+    """Test file parses as store content."""
+    content = """
+amd64:
+- install-type: store
+  name: k8s
+  channel: edge
+"""
+    mock_open().__enter__.return_value = io.StringIO(content)
+    mock_checkoutput().decode.return_value = "amd64"
+    args = snap._parse_management_arguments()
+    assert args == [
+        snap.SnapStoreArgument(name="k8s", channel="edge"),
+    ]
+
+
+@mock.patch("pathlib.Path.exists", mock.Mock(return_value=True))
+@mock.patch("pathlib.Path.open")
+@mock.patch("subprocess.check_output")
+def test_parse_valid_file(mock_checkoutput, mock_open):
+    """Test file parses as file content."""
+    content = """
+amd64:
+- install-type: file
+  name: k8s
+  filename: path/to/thing
+"""
+    mock_open().__enter__.return_value = io.StringIO(content)
+    mock_checkoutput().decode.return_value = "amd64"
+    args = snap._parse_management_arguments()
+    assert args == [
+        snap.SnapFileArgument(name="k8s", filename=Path("path/to/thing")),
+    ]
+
+
+@mock.patch("snap._parse_management_arguments")
+@mock.patch("snap.snap_lib.install_local")
+@mock.patch("snap.snap_lib.SnapCache")
+def test_management_installs_local(cache, install_local, args):
+    """Test installer uses local installer."""
+    cache.return_value.__getitem__.return_value = mock.MagicMock(spec=snap.snap_lib.Snap)
+    args.return_value = [
+        snap.SnapFileArgument(name="k8s", filename=Path("path/to/thing")),
+    ]
+    snap.management()
+    cache.called_once_with()
+    cache["k8s"].ensure.assert_not_called()
+    install_local.assert_called_once_with(filename=Path("path/to/thing"))
+
+
+@mock.patch("snap._parse_management_arguments")
+@mock.patch("snap.snap_lib.install_local")
+@mock.patch("snap.snap_lib.SnapCache")
+def test_management_installs_store(cache, install_local, args):
+    """Test installer uses store installer."""
+    cache.return_value.__getitem__.return_value = mock.MagicMock(spec=snap.snap_lib.Snap)
+    args.return_value = [
+        snap.SnapStoreArgument(name="k8s", channel="edge"),
+    ]
+    snap.management()
+    cache.called_once_with()
+    install_local.assert_not_called()
+    cache()["k8s"].ensure.assert_called_once_with(
+        state=snap.snap_lib.SnapState.Present, channel="edge"
+    )