Make root requirement check a decorator

Author: jwlodek

src/manage_iocs/commands.py

@@ -1,5 +1,4 @@
 import inspect
-import os
 import socket
 import sys
 from subprocess import PIPE, Popen
@@ -75,6 +74,7 @@ def report():
         )
 
 
+@utils.requires_root
 def disable(ioc: str):
     """Disable autostart for the given IOC."""
 
@@ -86,6 +86,7 @@ def disable(ioc: str):
     return ret
 
 
+@utils.requires_root
 def enable(ioc: str):
     """Enable autostart for the given IOC."""
 
@@ -139,6 +140,7 @@ def stopall():
     return ret
 
 
+@utils.requires_root
 def enableall():
     """Enable autostart for all IOCs on this host."""
 
@@ -149,6 +151,7 @@ def enableall():
     return ret
 
 
+@utils.requires_root
 def disableall():
     """Disable autostart for all IOCs on this host."""
 
@@ -170,12 +173,10 @@ def restart(ioc: str):
     return ret
 
 
+@utils.requires_root
 def uninstall(ioc: str):
     """Uninstall the given IOC."""
 
-    if not os.geteuid() == 0:
-        raise RuntimeError("You must be root to uninstall an IOC!")
-
     _, _, ret = utils.systemctl_passthrough("stop", ioc)
     if ret != 0:
         raise RuntimeError(f"Failed to stop IOC '{ioc}' before uninstalling!")
@@ -190,12 +191,10 @@ def uninstall(ioc: str):
     return ret
 
 
+@utils.requires_root
 def install(ioc: str):
     """Install the given IOC."""
 
-    if not os.geteuid() == 0:
-        raise RuntimeError("You must be root to install an IOC!")
-
     service_file = utils.SYSTEMD_SERVICE_PATH / f"softioc-{ioc}.service"
     ioc_config = utils.find_iocs()[ioc]
 

src/manage_iocs/utils.py

@@ -1,5 +1,6 @@
 import os
 import socket
+from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
 from subprocess import PIPE, Popen
@@ -107,3 +108,16 @@ def get_ioc_statuses(ioc_name: str) -> tuple[int, tuple[str, str]]:
     ret = ret + ret_enable
 
     return ret, (status, enabled.capitalize())
+
+
+def requires_root(func: Callable):
+    def wrapper(*args, **kwargs):
+        if os.geteuid() != 0:
+            raise PermissionError(f"Command {func.__name__} requires root privileges.")
+        return func(*args, **kwargs)
+
+    # Preserve the original function's name and docstring
+    wrapper.__doc__ = func.__doc__
+    wrapper.__name__ = func.__name__
+
+    return wrapper

tests/conftest.py

@@ -14,6 +14,11 @@ def all_manage_iocs_commands():
     return [obj for name, obj in inspect.getmembers(cmds) if inspect.isfunction(obj)]
 
 
+@pytest.fixture(autouse=True)
+def sim_running_as_root(monkeypatch):
+    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
+
+
 @pytest.fixture
 def sample_config_file_factory(tmp_path):
     def _simple_config_file(

tests/test_commands.py

@@ -45,26 +45,37 @@ def normalize_whitespace(s: str) -> str:
 
 
 @pytest.mark.parametrize(
-    "ioc_name, command, before_state, before_enabled, after_state, after_enabled",
+    "ioc_name, command, before_state, before_enabled, after_state, after_enabled, as_root",
     [
-        ("ioc1", cmds.stop, "Running", "Enabled", "Stopped", "Enabled"),
-        ("ioc3", cmds.stop, "Running", "Disabled", "Stopped", "Disabled"),
-        ("ioc4", cmds.stop, "Stopped", "Disabled", "Stopped", "Disabled"),
-        ("ioc1", cmds.start, "Running", "Enabled", "Running", "Enabled"),
-        ("ioc3", cmds.start, "Running", "Disabled", "Running", "Disabled"),
-        ("ioc4", cmds.start, "Stopped", "Disabled", "Running", "Disabled"),
-        ("ioc3", cmds.enable, "Running", "Disabled", "Running", "Enabled"),
-        ("ioc4", cmds.enable, "Stopped", "Disabled", "Stopped", "Enabled"),
-        ("ioc1", cmds.disable, "Running", "Enabled", "Running", "Disabled"),
-        ("ioc3", cmds.disable, "Running", "Disabled", "Running", "Disabled"),
-        ("ioc1", cmds.restart, "Running", "Enabled", "Running", "Enabled"),
-        ("ioc4", cmds.restart, "Stopped", "Disabled", "Running", "Disabled"),
-        ("ioc3", cmds.restart, "Running", "Disabled", "Running", "Disabled"),
+        ("ioc1", cmds.stop, "Running", "Enabled", "Stopped", "Enabled", False),
+        ("ioc3", cmds.stop, "Running", "Disabled", "Stopped", "Disabled", False),
+        ("ioc4", cmds.stop, "Stopped", "Disabled", "Stopped", "Disabled", False),
+        ("ioc1", cmds.start, "Running", "Enabled", "Running", "Enabled", False),
+        ("ioc3", cmds.start, "Running", "Disabled", "Running", "Disabled", False),
+        ("ioc4", cmds.start, "Stopped", "Disabled", "Running", "Disabled", False),
+        ("ioc3", cmds.enable, "Running", "Disabled", "Running", "Enabled", True),
+        ("ioc4", cmds.enable, "Stopped", "Disabled", "Stopped", "Enabled", True),
+        ("ioc1", cmds.disable, "Running", "Enabled", "Running", "Disabled", True),
+        ("ioc3", cmds.disable, "Running", "Disabled", "Running", "Disabled", True),
+        ("ioc1", cmds.restart, "Running", "Enabled", "Running", "Enabled", False),
+        ("ioc4", cmds.restart, "Stopped", "Disabled", "Running", "Disabled", False),
+        ("ioc3", cmds.restart, "Running", "Disabled", "Running", "Disabled", False),
     ],
 )
 def test_state_change_commands(
-    sample_iocs, ioc_name, command, before_state, before_enabled, after_state, after_enabled
+    sample_iocs,
+    ioc_name,
+    command,
+    before_state,
+    before_enabled,
+    after_state,
+    after_enabled,
+    as_root,
+    monkeypatch,
 ):
+    if not as_root:
+        monkeypatch.setattr(os, "geteuid", lambda: 1000)  # Mock as non-root user
+
     _, status = get_ioc_statuses(ioc_name)
     assert status == (before_state, before_enabled)
 
@@ -76,8 +87,6 @@ def test_state_change_commands(
 
 
 def test_install_new_ioc(sample_iocs, monkeypatch):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     assert "ioc2" not in find_installed_iocs()
 
     rc = cmds.install("ioc2")
@@ -89,32 +98,19 @@ def test_install_new_ioc(sample_iocs, monkeypatch):
     assert "ioc2" in find_installed_iocs()
 
 
-def test_install_ioc_not_root(sample_iocs, monkeypatch):
-    monkeypatch.setattr(os, "geteuid", lambda: 1000)  # Mock as non-root user
-
-    with pytest.raises(RuntimeError, match="You must be root to install an IOC!"):
-        cmds.install("ioc3")
-
-
 def test_install_ioc_wrong_host(sample_iocs, monkeypatch):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     with pytest.raises(RuntimeError, match="Cannot install IOC 'ioc1' on this host"):
         cmds.install("ioc1")
 
 
 def test_install_already_installed_ioc(sample_iocs, monkeypatch):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     assert "ioc3" in find_installed_iocs()
 
     with pytest.raises(RuntimeError, match="Failed to install IOC 'ioc3'!"):
         cmds.install("ioc3")
 
 
 def test_uninstall_ioc(sample_iocs, monkeypatch):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     assert "ioc5" in find_installed_iocs()
 
     rc = cmds.uninstall("ioc5")
@@ -123,11 +119,15 @@ def test_uninstall_ioc(sample_iocs, monkeypatch):
     assert "ioc5" not in find_installed_iocs()
 
 
-def test_uninstall_ioc_not_root(sample_iocs, monkeypatch):
+@pytest.mark.parametrize(
+    "command",
+    [cmds.enable, cmds.disable, cmds.enableall, cmds.disableall, cmds.install, cmds.uninstall],
+)
+def test_requires_root(sample_iocs, monkeypatch, command):
     monkeypatch.setattr(os, "geteuid", lambda: 1000)  # Mock as non-root user
 
-    with pytest.raises(RuntimeError, match="You must be root to uninstall an IOC!"):
-        cmds.uninstall("ioc1")
+    with pytest.raises(PermissionError, match="requires root privileges."):
+        command("ioc1")
 
 
 @pytest.mark.parametrize(
@@ -155,15 +155,21 @@ def test_state_change_all(sample_iocs, cmd, expected_state, expected_enabled):
             assert get_ioc_statuses(ioc.name)[1][1] == expected_enabled
 
 
-def test_attach(sample_iocs, monkeypatch, dummy_popen):
-    monkeypatch.setattr(os, "geteuid", lambda: 1000)  # Mock as non-root user
+@pytest.mark.parametrize("as_root", [True, False])
+def test_attach(sample_iocs, monkeypatch, dummy_popen, as_root):
+    if not as_root:
+        monkeypatch.setattr(os, "geteuid", lambda: 1000)  # Mock as non-root user
 
     ret = cmds.attach("ioc3")
 
     assert ret == ["telnet", "localhost", "3456"]
 
 
-def test_status(sample_iocs, capsys):
+@pytest.mark.parametrize("as_root", [True, False])
+def test_status(sample_iocs, capsys, monkeypatch, as_root):
+    if not as_root:
+        monkeypatch.setattr(os, "geteuid", lambda: 1000)  # Mock as non-root user
+
     rc = cmds.status()
     captured = capsys.readouterr()
     expected_output = """IOC Status Auto-Start
@@ -194,8 +200,6 @@ def normalize_whitespace(s: str) -> str:
     ],
 )
 def test_command_failures(sample_iocs, monkeypatch, cmd, expected_message):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     def failing_systemctl_passthrough(action: str, ioc: str) -> tuple[str, str, int]:
         return ("", "Simulated failure", 1)
 
@@ -214,8 +218,6 @@ def failing_systemctl_passthrough(action: str, ioc: str) -> tuple[str, str, int]
     ],
 )
 def test_uninstall_failures(sample_iocs, monkeypatch, failed_action, expected_message):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     def failing_systemctl_passthrough(action: str, ioc: str) -> tuple[str, str, int]:
         if action == failed_action:
             return ("", "Simulated failure", 1)
@@ -228,8 +230,6 @@ def failing_systemctl_passthrough(action: str, ioc: str) -> tuple[str, str, int]
 
 
 def test_fail_to_install_ioc_to_run_as_root(sample_iocs, monkeypatch, sample_config_file_factory):
-    monkeypatch.setattr(os, "geteuid", lambda: 0)  # Mock as root user
-
     sample_config_file_factory(name="ioc1", user="root")
     with pytest.raises(RuntimeError, match="Refusing to install IOC 'ioc1' to run as user 'root'!"):
         cmds.install("ioc1")