From 57dc31345f928cf84e90c63f4fa3884304cbbaab Mon Sep 17 00:00:00 2001
From: Daniel Perrefort <djperrefort@gmail.com>
Date: Thu, 7 Dec 2023 13:35:06 -0500
Subject: [PATCH] Updates tests for include_user_whitelist

---
 tests/utils/test_id_in_whitelist.py        | 27 -------------
 tests/utils/test_include_user_whitelist.py | 45 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 27 deletions(-)
 delete mode 100644 tests/utils/test_id_in_whitelist.py
 create mode 100644 tests/utils/test_include_user_whitelist.py

diff --git a/tests/utils/test_id_in_whitelist.py b/tests/utils/test_id_in_whitelist.py
deleted file mode 100644
index 114cdcd..0000000
--- a/tests/utils/test_id_in_whitelist.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from unittest import TestCase
-
-from shinigami.utils import id_in_whitelist
-
-
-class Whitelisting(TestCase):
-    """Tests for the ``id_in_whitelist`` function"""
-
-    def test_empty_whitelist(self) -> None:
-        """Test the return value is ``False`` for all ID values when the whitelist is empty"""
-
-        self.assertFalse(id_in_whitelist(0, []))
-        self.assertFalse(id_in_whitelist(123, []))
-
-    def test_whitelisted_by_id(self) -> None:
-        """Test return values for a whitelist of explicit ID values"""
-
-        whitelist = (123, 456, 789)
-        self.assertTrue(id_in_whitelist(456, whitelist))
-        self.assertFalse(id_in_whitelist(0, whitelist))
-
-    def test_whitelisted_by_id_range(self) -> None:
-        """Test return values for a whitelist of ID ranges"""
-
-        whitelist = (0, 1, 2, (100, 300))
-        self.assertTrue(id_in_whitelist(123, whitelist))
-        self.assertFalse(id_in_whitelist(301, whitelist))
diff --git a/tests/utils/test_include_user_whitelist.py b/tests/utils/test_include_user_whitelist.py
new file mode 100644
index 0000000..22c33b6
--- /dev/null
+++ b/tests/utils/test_include_user_whitelist.py
@@ -0,0 +1,45 @@
+"""Tests for the `include_user_whitelist` function."""
+
+from unittest import TestCase
+
+import pandas as pd
+
+from shinigami.utils import INIT_PROCESS_ID, include_user_whitelist
+
+
+class WhitelistedIDs(TestCase):
+    """Tests for the ``id_in_whitelist`` function"""
+
+    def setUp(self) -> None:
+        """Define a DataFrame with example process data
+
+        The returned DataFrame includes an init process, a mix of (non)orphaned
+        processes, and multiple user IDs.
+        """
+
+        self.testing_data = pd.DataFrame({
+            'PID': [INIT_PROCESS_ID, 100, 101, 102, 103],
+            'PPID': [0, INIT_PROCESS_ID, INIT_PROCESS_ID, 100, 200],
+            'PGID': [0, 1, 1, 2, 3],
+            'UID': [0, 123, 123, 456, 789],
+            'CMD': ['init', 'process_1', 'process_2', 'process_3', 'process_4']})
+
+    def test_empty_whitelist(self) -> None:
+        """Test the returned DataFrame is empty when the whitelist is empty"""
+
+        returned_df = include_user_whitelist(self.testing_data, [])
+        self.assertTrue(returned_df.empty)
+
+    def test_whitelisted_by_id(self) -> None:
+        """Test the returned UID values against a whitelist of explicit ID values"""
+
+        whitelist = (123, 456)
+        returned_df = include_user_whitelist(self.testing_data, whitelist)
+        self.assertTrue(set(returned_df['UID']).issubset(whitelist))
+
+    def test_whitelisted_by_id_range(self) -> None:
+        """Test the returned UID values against a whitelist of ID ranges"""
+
+        whitelist = (1, 2, (100, 500))
+        returned_df = include_user_whitelist(self.testing_data, whitelist)
+        self.assertCountEqual(returned_df['UID'].unique(), {123, 456})