feat: pin the grafana-agent snap installed by the charm

This:
* updates snap installations so that they pin their snaps to a specific version.  The reason being that we want each charm to ship with a specific version of a snap, and upgrades to happen explicitly with the charm rather than implicitly behind the scenes.  
* pins the snaps in this charm to the current latest/stable versions for each arch

src/charm.py

@@ -21,6 +21,7 @@
 from grafana_agent import METRICS_RULES_SRC_PATH, GrafanaAgentCharm
 from ops.main import main
 from ops.model import BlockedStatus, MaintenanceStatus, Relation
+from snap_management import SnapSpecError, install_ga_snap
 
 logger = logging.getLogger(__name__)
 
@@ -179,7 +180,6 @@ def __init__(self, *args):
         # we always listen to juju-info-joined events even though one of the two paths will be
         # at all effects unused.
         self._cos = COSAgentRequirer(self)
-        self.snap = snap.SnapCache()["grafana-agent"]
         self._tracing = TracingEndpointRequirer(self, protocols=["otlp_http"])
         self.framework.observe(
             self._cos.on.data_changed,  # pyright: ignore
@@ -194,6 +194,12 @@ def __init__(self, *args):
         self.framework.observe(self.on.stop, self._on_stop)
         self.framework.observe(self.on.remove, self._on_remove)
 
+    @property
+    def snap(self):
+        """Return the snap object for the Grafana Agent snap."""
+        # This is handled in a property to avoid calls to snapd until they're necessary.
+        return snap.SnapCache()["grafana-agent"]
+
     def _on_juju_info_joined(self, _event):
         """Update the config when Juju info is joined."""
         self._update_config()
@@ -222,11 +228,14 @@ def _on_cos_validation_error(self, event):
 
     def on_install(self, _event) -> None:
         """Install the Grafana Agent snap."""
-        # Check if Grafana Agent is installed
+        self._install()
+
+    def _install(self) -> None:
+        """Install/refresh the Grafana Agent snap."""
         self.unit.status = MaintenanceStatus("Installing grafana-agent snap")
         try:
-            self.snap.ensure(state=snap.SnapState.Latest)
-        except snap.SnapError as e:
+            install_ga_snap(classic=False)
+        except (snap.SnapError, SnapSpecError) as e:
             raise GrafanaAgentInstallError("Failed to install grafana-agent.") from e
 
     def _on_start(self, _event) -> None:
@@ -259,6 +268,12 @@ def _on_remove(self, _event) -> None:
         except snap.SnapError as e:
             raise GrafanaAgentInstallError("Failed to uninstall grafana-agent") from e
 
+    def _on_upgrade_charm(self, event):
+        """Upgrade the charm."""
+        # This is .observe()'d in the base class and thus not observed here
+        super()._on_upgrade_charm(event)
+        self._install()
+
     @property
     def is_k8s(self) -> bool:
         """Is this a k8s charm."""

src/snap_management.py

@@ -0,0 +1,82 @@
+#!/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.
+
+Modified from https://github.com/canonical/k8s-operator/blob/main/charms/worker/k8s/src/snap.py
+"""
+
+
+import logging
+import platform
+
+import charms.operator_libs_linux.v2.snap as snap_lib
+
+# Log messages can be retrieved using juju debug-log
+log = logging.getLogger(__name__)
+
+
+# Map of the grafana-agent snap revision to install for given architectures and strict mode.
+_grafana_agent_snap_name = "grafana-agent"
+_grafana_agent_snaps = {
+    # (confinement, arch): revision
+    ("strict", "amd64"): 16,
+    ("strict", "arm64"): 23,
+}
+
+
+class SnapSpecError(Exception):
+    """Custom exception type for errors related to the snap spec."""
+
+    pass
+
+
+def install_ga_snap(classic: bool):
+    """Looks up system details and installs the appropriate grafana-agent snap revision."""
+    arch = get_system_arch()
+    confinement = "classic" if classic else "strict"
+    try:
+        revision = str(_grafana_agent_snaps[(confinement, arch)])
+    except KeyError as e:
+        raise SnapSpecError(
+            f"Snap spec not found for arch={arch} and confinement={confinement}"
+        ) from e
+    _install_snap(name=_grafana_agent_snap_name, revision=revision, classic=classic)
+
+
+def _install_snap(
+    name: str,
+    revision: str,
+    classic: bool = False,
+):
+    """Install and pin the given snap revision.
+
+    The revision will be held, i.e. it won't be automatically updated any time a new revision is released.
+    """
+    cache = snap_lib.SnapCache()
+    snap = cache[name]
+    log.info(
+        f"Ensuring {name} snap is installed at revision={revision}"
+        f" with classic confinement={classic}"
+    )
+    snap.ensure(state=snap_lib.SnapState.Present, revision=revision, classic=classic)
+    snap.hold()
+
+
+def get_system_arch() -> str:
+    """Returns the architecture of this machine, mapping some values to amd64 or arm64.
+
+    If platform is x86_64 or amd64, it returns amd64.
+    If platform is aarch64, arm64, armv8b, or armv8l, it returns arm64.
+    """
+    arch = platform.processor()
+    if arch in ["x86_64", "amd64"]:
+        arch = "amd64"
+    elif arch in ["aarch64", "arm64", "armv8b", "armv8l"]:
+        arch = "arm64"
+    # else: keep arch as is
+    return arch

tests/scenario/test_machine_charm/conftest.py

@@ -1,8 +1,16 @@
 # Copyright 2022 Canonical Ltd.
 # See LICENSE file for licensing details.
+from unittest.mock import patch
+
 import pytest
 
 
 @pytest.fixture
 def placeholder_cfg_path(tmp_path):
     return tmp_path / "foo.yaml"
+
+
+@pytest.fixture()
+def mock_config_path(placeholder_cfg_path):
+    with patch("grafana_agent.CONFIG_PATH", placeholder_cfg_path):
+        yield

tests/scenario/test_machine_charm/test_alert_labels.py

@@ -2,13 +2,28 @@
 # See LICENSE file for licensing details.
 
 import json
+from unittest.mock import PropertyMock, patch
 
 import charm
+import pytest
 from scenario import Context, PeerRelation, Relation, State, SubordinateRelation
 
 from tests.scenario.helpers import get_charm_meta
 
 
+@pytest.fixture(autouse=True)
+def use_mock_config_path(mock_config_path):
+    # Use the common mock_config_path fixture from conftest.py
+    yield
+
+
+@pytest.fixture(autouse=True)
+def mock_snap():
+    """Mock the charm's snap property so we don't access the host."""
+    with patch("charm.GrafanaAgentMachineCharm.snap", new_callable=PropertyMock):
+        yield
+
+
 def test_metrics_alert_rule_labels(vroot):
     """Check that metrics alert rules are labeled with principal topology."""
     cos_agent_primary_data = {
@@ -105,6 +120,7 @@ def test_metrics_alert_rule_labels(vroot):
             PeerRelation("peers"),
         ],
     )
+
     state_0 = context.run(event=cos_agent_primary_relation.changed_event, state=state)
     (vroot / "metadata.yaml").unlink(missing_ok=True)
     (vroot / "config.yaml").unlink(missing_ok=True)

tests/scenario/test_machine_charm/test_multiple_subordinates.py

@@ -2,13 +2,28 @@
 # See LICENSE file for licensing details.
 
 import json
+from unittest.mock import PropertyMock, patch
 
 import charm
+import pytest
 from scenario import Context, PeerRelation, State, SubordinateRelation
 
 from tests.scenario.helpers import get_charm_meta
 
 
+@pytest.fixture(autouse=True)
+def use_mock_config_path(mock_config_path):
+    # Use the common mock_config_path fixture from conftest.py
+    yield
+
+
+@pytest.fixture(autouse=True)
+def mock_snap():
+    """Mock the charm's snap property so we don't access the host."""
+    with patch("charm.GrafanaAgentMachineCharm.snap", new_callable=PropertyMock):
+        yield
+
+
 def test_juju_info_and_cos_agent(vroot):
     def post_event(charm: charm.GrafanaAgentMachineCharm):
         assert len(charm._cos.dashboards) == 1

tests/scenario/test_setup_statuses.py

@@ -2,7 +2,7 @@
 # See LICENSE file for licensing details.
 import dataclasses
 from typing import Type
-from unittest.mock import patch
+from unittest.mock import PropertyMock, patch
 
 import charm
 import grafana_agent
@@ -14,7 +14,7 @@
 from tests.scenario.helpers import get_charm_meta
 
 
-@pytest.fixture(params=["k8s", "lxd"])
+@pytest.fixture(params=["lxd"])
 def substrate(request):
     return request.param
 
@@ -46,7 +46,21 @@ def patch_all(substrate, mock_cfg_path):
         yield
 
 
-def test_install(charm_type, substrate, vroot):
+@pytest.fixture(autouse=True)
+def mock_snap():
+    """Mock the charm's snap property so we don't access the host."""
+    with patch(
+        "charm.GrafanaAgentMachineCharm.snap", new_callable=PropertyMock
+    ) as mocked_property:
+        mock_snap = mocked_property.return_value
+        # Mock the .present property of the snap object so the start event doesn't try to configure
+        # anything
+        mock_snap.present = False
+        yield
+
+
+@patch("charm.SnapManifest._get_system_arch", return_value="amd64")
+def test_install(_mock_manifest_get_system_arch, charm_type, substrate, vroot):
     context = Context(
         charm_type,
         meta=get_charm_meta(charm_type),

tests/unit/test_relation_status.py

@@ -26,6 +26,10 @@ def setUp(self, *unused):
         self.mock_snap = patcher.start()
         self.addCleanup(patcher.stop)
 
+        patcher = patch.object(GrafanaAgentCharm, "_install")
+        self.mock_install = patcher.start()
+        self.addCleanup(patcher.stop)
+
         self.harness = Harness(GrafanaAgentCharm)
         self.harness.set_model_name(self.__class__.__name__)
 

tests/unit/test_update_status.py

@@ -25,6 +25,10 @@ def setUp(self, *unused):
         self.mock_snap = patcher.start()
         self.addCleanup(patcher.stop)
 
+        patcher = patch.object(GrafanaAgentCharm, "_install")
+        self.mock_install = patcher.start()
+        self.addCleanup(patcher.stop)
+
         self.harness = Harness(GrafanaAgentCharm)
         self.harness.set_model_name(self.__class__.__name__)