From 60f0c13ad903692467812421657514c408bb272a Mon Sep 17 00:00:00 2001
From: "Jason C. Nucciarone" <nuccitheboss@ubuntu.com>
Date: Tue, 2 Jul 2024 18:15:08 -0400
Subject: [PATCH] refactor: use `mungectl generate` rather than generating own
 bytes

Signed-off-by: Jason C. Nucciarone <nuccitheboss@ubuntu.com>
---
 src/slurmhelpers/models.py | 10 ++++++----
 tests/unit/conftest.py     |  3 +--
 tests/unit/test_hooks.py   |  1 +
 tests/unit/test_models.py  | 32 +++++++++++++++++---------------
 4 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/src/slurmhelpers/models.py b/src/slurmhelpers/models.py
index 91f99da..4be2a68 100644
--- a/src/slurmhelpers/models.py
+++ b/src/slurmhelpers/models.py
@@ -16,7 +16,7 @@
 
 import base64
 import logging
-import os
+import subprocess
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
 
@@ -165,10 +165,12 @@ def generate_key(self) -> None:
         Replicates the daemon autostart feature from the  munge Debian package.
         """
         logging.info("Generating new secret file for service `munged`.")
-        if not self.secret_file.exists():
-            self.secret_file.touch(0o600)
+        try:
+            subprocess.check_output(["mungectl", "generate"])
+        except subprocess.CalledProcessError:
+            logging.fatal("Failed to generate a new munge key")
+            raise
 
-        self.secret_file.write_bytes(os.urandom(1024))
         self._needs_restart(["munged"])
 
     def update_config(self, config: Dict[str, str]) -> None:
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
index 0d3cc3f..41c4a06 100644
--- a/tests/unit/conftest.py
+++ b/tests/unit/conftest.py
@@ -18,11 +18,10 @@
 from unittest.mock import MagicMock, PropertyMock
 
 import pytest
+from slurmhelpers.models import Munge, Slurm, Slurmd, Slurmdbd, Slurmrestd
 from snaphelpers import Snap, SnapConfig, SnapConfigOptions, SnapServices
 from snaphelpers._ctl import ServiceInfo
 
-from slurmhelpers.models import Munge, Slurm, Slurmd, Slurmdbd, Slurmrestd
-
 
 @pytest.fixture
 def env():
diff --git a/tests/unit/test_hooks.py b/tests/unit/test_hooks.py
index 0802bb7..6409adc 100644
--- a/tests/unit/test_hooks.py
+++ b/tests/unit/test_hooks.py
@@ -80,6 +80,7 @@ class TestHooks:
 
     def test_install_hook(self, mocker, snap) -> None:
         """Test `install` hook."""
+        mocker.patch("subprocess.check_output")
         mocker.patch("pathlib.Path.chmod")
         mocker.patch("pathlib.Path.mkdir")
         mocker.patch("pathlib.Path.touch")
diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py
index 4d49ea1..aa39863 100644
--- a/tests/unit/test_models.py
+++ b/tests/unit/test_models.py
@@ -15,7 +15,16 @@
 
 """Test models that wrap the configuration for the bundled daemons."""
 
+import subprocess
+
 import pytest
+from slurmhelpers.models import (
+    _process_down_nodes,
+    _process_frontend_nodes,
+    _process_node_sets,
+    _process_nodes,
+    _process_partitions,
+)
 from slurmutils.models import (
     DownNodes,
     DownNodesList,
@@ -31,14 +40,6 @@
     SlurmdbdConfig,
 )
 
-from slurmhelpers.models import (
-    _process_down_nodes,
-    _process_frontend_nodes,
-    _process_node_sets,
-    _process_nodes,
-    _process_partitions,
-)
-
 
 class TestBaseModel:
     """Test the `_BaseModel` parent class for data models."""
@@ -114,15 +115,16 @@ def test_max_thread_count(self, mocker, munge) -> None:
     def test_generate_key(self, mocker, munge) -> None:
         """Test `generate_key` method."""
         # Generate key when `munge.key` file does not yet exist.
-        mocker.patch("pathlib.Path.exists", return_value=False)
-        mocker.patch("pathlib.Path.touch")
-        mocker.patch("pathlib.Path.write_bytes")
+        mocker.patch("subprocess.check_output")
         munge.generate_key()
 
-        # Generate key when `munge.key` file exists.
-        mocker.patch("pathlib.Path.exists", return_value=True)
-        mocker.patch("pathlib.Path.write_bytes")
-        munge.generate_key()
+        # Fail to generate new `munge.key`.
+        mocker.patch(
+            "subprocess.check_output",
+            side_effect=subprocess.CalledProcessError(0, ["mungectl", "generate"]),
+        )
+        with pytest.raises(subprocess.CalledProcessError):
+            munge.generate_key()
 
     def test_update_config(self, mocker, munge) -> None:
         """Test `update_config` method."""