Skip to content

Commit

Permalink
Updates tests for include_user_whitelist
Browse files Browse the repository at this point in the history
  • Loading branch information
djperrefort committed Dec 7, 2023
1 parent ffa34c0 commit 57dc313
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 27 deletions.
27 changes: 0 additions & 27 deletions tests/utils/test_id_in_whitelist.py

This file was deleted.

45 changes: 45 additions & 0 deletions tests/utils/test_include_user_whitelist.py
Original file line number Diff line number Diff line change
@@ -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})

0 comments on commit 57dc313

Please sign in to comment.