db: add a unit test for update_max_threat_level()

stratosphereips · Dec 1, 2023 · 05626c8 · 05626c8
1 parent 6be3c9f
commit 05626c8
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/slips_files/core/database/database_manager.py b/slips_files/core/database/database_manager.py
@@ -831,6 +831,12 @@ def execute_query(self, *args, **kwargs):
     def get_pid_of(self, *args, **kwargs):
         return self.rdb.get_pid_of(*args, **kwargs)
 
+    def set_max_threat_level(self, *args, **kwargs):
+        return self.rdb.set_max_threat_level(*args, **kwargs)
+
+    def update_max_threat_level(self, *args, **kwargs):
+        return self.rdb.update_max_threat_level(*args, **kwargs)
+
     def get_name_of_module_at(self, *args, **kwargs):
         return self.rdb.get_name_of_module_at(*args, **kwargs)
 

diff --git a/tests/test_database.py b/tests/test_database.py
@@ -1,5 +1,6 @@
 from slips_files.common.slips_utils import utils
 from slips_files.core.flows.zeek import Conn
+from slips_files.common.slips_utils import utils
 from tests.module_factory import ModuleFactory
 import redis
 import os
@@ -221,3 +222,19 @@ def test_get_the_other_ip_version():
 def test_add_tuple(tupleid: str, symbol, expected_direction, role, flow):
     db.add_tuple(profileid, twid, tupleid, symbol, role, flow)
     assert symbol[0] in db.r.hget(f'profile_{flow.saddr}_{twid}', expected_direction)
+
+
+@pytest.mark.parametrize(
+    'max_threat_level, cur_threat_level, expected_max',
+    [
+        ('info', 'info', utils.threat_levels['info']),
+        ('critical', 'info', utils.threat_levels['critical']),
+        ('high', 'critical', utils.threat_levels['critical']),
+    ],
+)
+def test_update_max_threat_level(
+        max_threat_level, cur_threat_level, expected_max
+    ):
+    db.set_max_threat_level(profileid, max_threat_level)
+    assert db.update_max_threat_level(
+        profileid, cur_threat_level) == expected_max