From 54bcdada066a5c3a378bbfe13086e30b358723ec Mon Sep 17 00:00:00 2001
From: Ray Chan <72372217+chanchiwai-ray@users.noreply.github.com>
Date: Thu, 29 Feb 2024 09:58:24 +0800
Subject: [PATCH] Enable / disable exporter service when start / stop the
 service. (#174)

---
 src/charm.py                |  6 ++++--
 src/service.py              | 10 ++++++++++
 tests/unit/test_exporter.py |  6 ++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index d28ddcad..32bd7b35 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -234,15 +234,17 @@ def _on_cos_agent_relation_joined(self, event: EventBase) -> None:
             )
             event.defer()
             return
+        self.exporter.enable()
         self.exporter.start()
-        logger.info("Start exporter service")
+        logger.info("Start and enable exporter service")
         self._on_update_status(event)
 
     def _on_cos_agent_relation_departed(self, event: EventBase) -> None:
         """Remove the exporter when relation departed."""
         if self._stored.exporter_installed:  # type: ignore[truthy-function]
             self.exporter.stop()
-            logger.info("Stop exporter service")
+            self.exporter.disable()
+            logger.info("Stop and disable exporter service")
         self._on_update_status(event)
 
     def _get_redfish_creds(self) -> Dict[str, str]:
diff --git a/src/service.py b/src/service.py
index f0caa8ac..d0fa65ae 100644
--- a/src/service.py
+++ b/src/service.py
@@ -159,6 +159,16 @@ def restart(self) -> None:
         """Restart the exporter daemon."""
         systemd.service_restart(EXPORTER_NAME)
 
+    @check_installed
+    def enable(self) -> None:
+        """Enable the exporter service."""
+        systemd.service_enable(EXPORTER_NAME)
+
+    @check_installed
+    def disable(self) -> None:
+        """Restart the exporter service."""
+        systemd.service_disable(EXPORTER_NAME)
+
     @check_installed
     def check_active(self) -> bool:
         """Check if the exporter is active or not."""
diff --git a/tests/unit/test_exporter.py b/tests/unit/test_exporter.py
index 41ee650f..dcf92170 100644
--- a/tests/unit/test_exporter.py
+++ b/tests/unit/test_exporter.py
@@ -120,6 +120,7 @@ def test_20_start_okay(self, mock_service_installed):
         self.harness.charm._stored.exporter_installed = True
         self.harness.add_relation_unit(rid, "grafana-agent/0")
         self.mock_systemd.service_start.assert_called_once()
+        self.mock_systemd.service_enable.assert_called_once()
 
     @mock.patch.object(pathlib.Path, "exists", return_value=False)
     def test_21_start_failed(self, mock_service_not_installed):
@@ -128,6 +129,7 @@ def test_21_start_failed(self, mock_service_not_installed):
         self.harness.begin()
         self.harness.add_relation_unit(rid, "grafana-agent/0")
         self.mock_systemd.service_start.assert_not_called()
+        self.mock_systemd.service_enable.assert_not_called()
 
     @mock.patch.object(pathlib.Path, "exists", return_value=True)
     def test_22_start_defer_resource_not_ready(self, mock_service_installed):
@@ -138,6 +140,7 @@ def test_22_start_defer_resource_not_ready(self, mock_service_installed):
         self.harness.charm._stored.exporter_installed = True
         self.harness.add_relation_unit(rid, "grafana-agent/0")
         self.mock_systemd.service_start.assert_not_called()
+        self.mock_systemd.service_enable.assert_not_called()
 
     @mock.patch.object(pathlib.Path, "exists", return_value=True)
     def test_23_start_defer_exporter_not_ready(self, mock_service_installed):
@@ -148,6 +151,7 @@ def test_23_start_defer_exporter_not_ready(self, mock_service_installed):
         self.harness.charm._stored.exporter_installed = False
         self.harness.add_relation_unit(rid, "grafana-agent/0")
         self.mock_systemd.service_start.assert_not_called()
+        self.mock_systemd.service_enable.assert_not_called()
 
     @mock.patch.object(pathlib.Path, "exists", return_value=True)
     def test_30_stop_okay(self, mock_service_installed):
@@ -159,6 +163,7 @@ def test_30_stop_okay(self, mock_service_installed):
         self.harness.add_relation_unit(rid, "grafana-agent/0")
         self.harness.remove_relation_unit(rid, "grafana-agent/0")
         self.mock_systemd.service_stop.assert_called_once()
+        self.mock_systemd.service_disable.assert_called_once()
 
     @mock.patch.object(pathlib.Path, "exists", return_value=False)
     def test_31_stop_failed(self, mock_service_not_installed):
@@ -170,6 +175,7 @@ def test_31_stop_failed(self, mock_service_not_installed):
         self.harness.add_relation_unit(rid, "grafana-agent/0")
         self.harness.remove_relation_unit(rid, "grafana-agent/0")
         self.mock_systemd.service_stop.assert_not_called()
+        self.mock_systemd.service_disable.assert_not_called()
 
     @parameterized.expand(
         [