From 77af4c2da640e51976f4ae1ed48c650f699e3e43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Juli=C3=A1n=20Espina?=
 <julian.espina@canonical.com>
Date: Thu, 11 Jul 2024 16:44:06 -0600
Subject: [PATCH] feat: add method to check if a service is active

---
 lib/charms/hpc_libs/v0/slurm_ops.py         | 14 ++++++++++++++
 tests/integration/slurm_ops/test_manager.py | 10 ++++++++--
 tests/unit/test_slurm_ops.py                | 10 ++++++++++
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/lib/charms/hpc_libs/v0/slurm_ops.py b/lib/charms/hpc_libs/v0/slurm_ops.py
index d49a38a..c34b38f 100644
--- a/lib/charms/hpc_libs/v0/slurm_ops.py
+++ b/lib/charms/hpc_libs/v0/slurm_ops.py
@@ -208,6 +208,20 @@ def restart(self) -> None:
         """Restart service."""
         _snap("restart", f"slurm.{self._service.value}")
 
+    def active(self) -> bool:
+        """Return True if the service is active."""
+        info = yaml.safe_load(_snap("info", "slurm"))
+        services = info["services"]
+
+        if (state := services.get(f"slurm.{self._service.value}")) is None:
+            _logger.warn(f"Tried to query for unknown service slurm.{self._service.value}")
+            return False
+
+        # We could also parse this into type, startup and status, but I don't think
+        # it's worth the additional complexity.
+        # We don't do `active in state` because the word "active" is also part of "inactive" :)
+        return "inactive" not in state
+
 
 class ConfigurationManager:
     """Control configuration of a Slurm component."""
diff --git a/tests/integration/slurm_ops/test_manager.py b/tests/integration/slurm_ops/test_manager.py
index 78ac8ca..bc77dfc 100644
--- a/tests/integration/slurm_ops/test_manager.py
+++ b/tests/integration/slurm_ops/test_manager.py
@@ -19,7 +19,6 @@ def slurmctld() -> SlurmManagerBase:
 def test_install(slurmctld: SlurmManagerBase) -> None:
     """Install Slurm using the manager."""
     slurm.install()
-    slurmctld.enable()
     slurmctld.munge.generate_key()
 
     with open("/var/snap/slurm/common/etc/munge/munge.key", "rb") as f:
@@ -40,7 +39,7 @@ def test_rotate_key(slurmctld: SlurmManagerBase) -> None:
 @pytest.mark.order(3)
 def test_slurm_config(slurmctld: SlurmManagerBase) -> None:
     """Test that the slurm config can be changed."""
-    slurmctld.config.set({"cluster-name": "test-cluster"})
+    slurmctld.config.set({"slurmctld-host": "test-slurm-ops", "cluster-name": "test-cluster"})
     value = slurmctld.config.get("cluster-name")
     assert value == "test-cluster"
 
@@ -57,6 +56,13 @@ def test_slurm_config(slurmctld: SlurmManagerBase) -> None:
 
 
 @pytest.mark.order(4)
+def test_enable_service(slurmctld: SlurmManagerBase) -> None:
+    """Test that the slurmctl daemon can be enabled."""
+    slurmctld.enable()
+    assert slurmctld.active()
+
+
+@pytest.mark.order(5)
 def test_version() -> None:
     """Test that the Slurm manager can report its version."""
     version = slurm.version()
diff --git a/tests/unit/test_slurm_ops.py b/tests/unit/test_slurm_ops.py
index 2d4b1c5..89b5d81 100644
--- a/tests/unit/test_slurm_ops.py
+++ b/tests/unit/test_slurm_ops.py
@@ -27,8 +27,13 @@
     - slurm.command1
     - slurm.command2
 services:
+    slurm.logrotate:                 oneshot, enabled, inactive
     slurm.munged:                    simple, enabled, active
+    slurm.slurm-prometheus-exporter: simple, disabled, inactive
     slurm.slurmctld:                 simple, disabled, active
+    slurm.slurmd:                    simple, enabled, active
+    slurm.slurmdbd:                  simple, disabled, active
+    slurm.slurmrestd:                simple, disabled, active
 channels:
     latest/stable:    –
     latest/candidate: 23.11.7 2024-06-26 (460) 114MB classic
@@ -101,6 +106,11 @@ def test_restart(self, subcmd, *_) -> None:
         args = subcmd.call_args[0][0]
         self.assertEqual(args, ["snap", "restart", f"slurm.{self.manager._service.value}"])
 
+    def test_active(self, subcmd, *_) -> None:
+        """Test that the manager can detect that a service is active."""
+        subcmd.return_value = SLURM_INFO.encode()
+        self.assertTrue(self.manager.active())
+
     def test_get_options(self, subcmd) -> None:
         """Test that the manager correctly collects all requested configuration options."""
         subcmd.return_value = '{"%(name)s.key1": "value1", "%(name)s.key2": "value2"}' % {