From d5ed16f6c947bc69807d21027749bc0c310b8a98 Mon Sep 17 00:00:00 2001
From: Arif Ali <arif.ali@canonical.com>
Date: Sat, 22 Jun 2024 00:34:51 +0100
Subject: [PATCH] [common] Standardise snap package checks

* Add function is_snap_installed, so that this can be used more widely
* This also reduces duplicate code in plugins.
* Add is_snap variable, that can be used if any of the packages is a
  snap; typically it's one of the other.

Signed-off-by: Arif Ali <arif.ali@canonical.com>
---
 sos/report/plugins/__init__.py | 18 ++++++++++++++++++
 sos/report/plugins/grafana.py  |  8 --------
 sos/report/plugins/kafka.py    |  8 --------
 sos/report/plugins/lxd.py      |  4 +---
 sos/report/plugins/maas.py     |  9 ---------
 5 files changed, 19 insertions(+), 28 deletions(-)

diff --git a/sos/report/plugins/__init__.py b/sos/report/plugins/__init__.py
index 191a96def9..07df31871b 100644
--- a/sos/report/plugins/__init__.py
+++ b/sos/report/plugins/__init__.py
@@ -541,6 +541,7 @@ class Plugin():
     cmdtags = {}
     filetags = {}
     option_list = []
+    is_snap = False
 
     # Default predicates
     predicate = None
@@ -581,6 +582,11 @@ def __init__(self, commons):
             opt.plugin = self.name()
             self.options[opt.name] = opt
 
+        # Check if any of the packages tuple is a snap
+        self.is_snap = any(
+            self.is_snap_installed(pkg) for pkg in list(self.packages)
+        )
+
         # Initialise the default --dry-run predicate
         self.set_predicate(SoSPredicate(self))
 
@@ -999,6 +1005,18 @@ def is_installed(self, package_name):
             len(self.policy.package_manager.all_pkgs_by_name(package_name)) > 0
         )
 
+    def is_snap_installed(self, package_name):
+        """Is the snap package $package_name installed?
+
+        :param package_name:    The name of the package to check
+        :type package_name:     ``str``
+
+        :returns: ``True`` if the snap package is installed, else ``False``
+        :rtype: ``bool``
+        """
+        pkg = self.policy.package_manager.pkg_by_name(package_name)
+        return pkg is not None and pkg['pkg_manager'] == 'snap'
+
     def is_service(self, name):
         """Does the service $name exist on the system?
 
diff --git a/sos/report/plugins/grafana.py b/sos/report/plugins/grafana.py
index 33d0b5587a..2b0ac44d8d 100644
--- a/sos/report/plugins/grafana.py
+++ b/sos/report/plugins/grafana.py
@@ -18,16 +18,8 @@ class Grafana(Plugin, IndependentPlugin):
     profiles = ('services', 'openstack', 'openstack_controller')
 
     packages = ('grafana',)
-    is_snap = False
-
-    def _is_snap_installed(self):
-        grafana_pkg = self.policy.package_manager.pkg_by_name('grafana')
-        if grafana_pkg:
-            return grafana_pkg['pkg_manager'] == 'snap'
-        return False
 
     def setup(self):
-        self.is_snap = self._is_snap_installed()
         if self.is_snap:
             grafana_cli = "grafana.grafana-cli"
             log_path = "/var/snap/grafana/common/data/log/"
diff --git a/sos/report/plugins/kafka.py b/sos/report/plugins/kafka.py
index 409210de8c..2a0a3582f8 100644
--- a/sos/report/plugins/kafka.py
+++ b/sos/report/plugins/kafka.py
@@ -23,16 +23,8 @@ class Kafka(Plugin, UbuntuPlugin):
     profiles = ('services',)
     packages = ('charmed-kafka',)
     services = ('kafka',)
-    is_snap = False
-
-    def _is_snap_installed(self):
-        kafka_pkg = self.policy.package_manager.pkg_by_name('charmed-kafka')
-        if kafka_pkg:
-            return kafka_pkg['pkg_manager'] == 'snap'
-        return False
 
     def setup(self):
-        self.is_snap = self._is_snap_installed()
         log_file_pattern = "*.log*" if self.get_option("all_logs") else "*.log"
 
         if self.is_snap:
diff --git a/sos/report/plugins/lxd.py b/sos/report/plugins/lxd.py
index 6df20dffd8..29f89dd487 100644
--- a/sos/report/plugins/lxd.py
+++ b/sos/report/plugins/lxd.py
@@ -21,9 +21,7 @@ class LXD(Plugin, UbuntuPlugin):
     services = ('snap.lxd.daemon', 'snap.lxd.activate')
 
     def setup(self):
-
-        lxd_pkg = self.policy.package_manager.pkg_by_name('lxd')
-        if lxd_pkg and lxd_pkg['pkg_manager'] == 'snap':
+        if self.is_snap:
 
             lxd_pred = SoSPredicate(self, services=['snap.lxd.daemon'],
                                     required={'services': 'all'})
diff --git a/sos/report/plugins/maas.py b/sos/report/plugins/maas.py
index 92dbe379e4..eb758d9742 100644
--- a/sos/report/plugins/maas.py
+++ b/sos/report/plugins/maas.py
@@ -47,8 +47,6 @@ class Maas(Plugin, UbuntuPlugin):
                   desc='Credentials, or the API key')
     ]
 
-    is_snap = False
-
     def _has_login_options(self):
         return self.get_option("url") and self.get_option("credentials") \
             and self.get_option("profile-name")
@@ -61,14 +59,7 @@ def _remote_api_login(self):
 
         return ret['status'] == 0
 
-    def _is_snap_installed(self):
-        maas_pkg = self.policy.package_manager.pkg_by_name('maas')
-        if maas_pkg:
-            return maas_pkg['pkg_manager'] == 'snap'
-        return False
-
     def setup(self):
-        self.is_snap = self._is_snap_installed()
         if self.is_snap:
             self.add_cmd_output([
                 'snap info maas',