Adds tests for exclude_active_slurm_users

shinigami/utils.py

@@ -120,8 +120,8 @@ def exclude_active_slurm_users(df: pd.DataFrame) -> pd.DataFrame:
         A copy of the given DataFrame
     """
 
-    is_slurm = df['CMD'].str.contains('slurmd').any()
-    slurm_uids = is_slurm['UID']
+    is_slurm = df['CMD'].str.contains('slurmd')
+    slurm_uids = df['UID'][is_slurm].unique()
     return df[~df['UID'].isin(slurm_uids)]
 
 
@@ -152,10 +152,10 @@ async def terminate_errant_processes(
         process_df = include_user_whitelist(process_df, uid_whitelist)
         process_df = exclude_active_slurm_users(process_df)
 
-        for _, row in process_df.iterrows():
+        for _, row in process_df.iterrows():  # pragma: nocover
             logging.info(f'[{node}] Marking for termination {dict(row)}')
 
-        if process_df.empty:
+        if process_df.empty:  # pragma: nocover
             logging.info(f'[{node}] no processes found')
 
         elif not debug:

tests/utils/test_exclude_active_slurm_users.py

@@ -0,0 +1,41 @@
+"""Tests for the `utils.exclude_active_slurm_users` function"""
+
+import unittest
+
+import pandas as pd
+
+from shinigami.utils import exclude_active_slurm_users
+
+
+class ExcludeSlurmUsers(unittest.TestCase):
+    """Test the identification of slurm users from a DataFrame of process data"""
+
+    def test_exclude_slurm_users(self) -> None:
+        input_df = pd.DataFrame({
+            'UID': [1001, 1002, 1003, 1004, 1005],
+            'CMD': ['process_1', 'slurmd', 'process_3', 'process_4', 'process_5']})
+
+        expected_df = input_df.loc[[0, 2, 3, 4]]
+        returned_df = exclude_active_slurm_users(input_df)
+
+        pd.testing.assert_frame_equal(returned_df, expected_df)
+        self.assertIsNot(returned_df, input_df)
+
+    def test_no_slurm_users(self) -> None:
+        input_df = pd.DataFrame({
+            'UID': [1001, 1002, 1003, 1004, 1005],
+            'CMD': ['process_1', 'process_2', 'process_3', 'process_4', 'process_5']})
+
+        returned_df = exclude_active_slurm_users(input_df)
+
+        pd.testing.assert_frame_equal(returned_df, input_df)
+        self.assertIsNot(returned_df, input_df)
+
+    def test_all_slurm_users(self) -> None:
+        input_df = pd.DataFrame({
+            'UID': [1001, 1002, 1003, 1004, 1005],
+            'CMD': ['slurmd', 'slurmd', 'slurmd', 'slurmd', 'slurmd']})
+
+        returned_df = exclude_active_slurm_users(input_df)
+        self.assertTrue(returned_df.empty)
+        self.assertIsNot(returned_df, input_df)

tests/utils/test_get_nodes.py

@@ -1,4 +1,4 @@
-"""Tests for the `utils` module."""
+"""Tests for the `utils.get_nodes` module."""
 
 import subprocess
 from unittest import TestCase, skipIf

tests/utils/test_include_orphaned_processes.py

@@ -1,4 +1,4 @@
-"""Tests for the `include_orphaned_processes` function."""
+"""Tests for the `utils.include_orphaned_processes` function."""
 
 import unittest
 

tests/utils/test_include_user_whitelist.py

@@ -1,4 +1,4 @@
-"""Tests for the `include_user_whitelist` function."""
+"""Tests for the `utils.include_user_whitelist` function."""
 
 from unittest import TestCase