Add test cases for stats

Author: bitterpanda63

aikido_zen/storage/statistics/__init__.py

@@ -10,8 +10,13 @@ class Statistics:
 
     def __init__(self):
         self.total_hits = 0
+
         self.attacks_detected = 0
         self.attacks_blocked = 0
+
+        self.attack_waves_detected = 0
+        self.attack_waves_blocked = 0
+
         self.rate_limited_hits = 0
         self.started_at = t.get_unixtime_ms()
         self.operations = Operations()
@@ -80,6 +85,8 @@ def empty(self):
             return False
         if self.attacks_detected > 0:
             return False
+        if self.attack_waves_detected > 0:
+            return False
         if len(self.operations) > 0:
             return False
         return True

aikido_zen/storage/statistics/init_test.py

@@ -18,6 +18,8 @@ def test_initialization(monkeypatch):
     assert stats.total_hits == 0
     assert stats.attacks_detected == 0
     assert stats.attacks_blocked == 0
+    assert stats.attack_waves_detected == 0
+    assert stats.attack_waves_blocked == 0
     assert stats.started_at == mock_time
     assert isinstance(stats.operations, Operations)
 
@@ -33,12 +35,15 @@ def test_clear(monkeypatch):
     stats.total_hits = 10
     stats.attacks_detected = 5
     stats.attacks_blocked = 3
+    stats.on_detected_attack_wave(blocked=True)
     stats.operations.register_call("test", "sql_op")
     stats.clear()
 
     assert stats.total_hits == 0
     assert stats.attacks_detected == 0
     assert stats.attacks_blocked == 0
+    assert stats.attack_waves_blocked == 0
+    assert stats.attack_waves_detected == 0
     assert stats.started_at == mock_time
     assert stats.operations == {}
 
@@ -59,6 +64,16 @@ def test_on_detected_attack(stats):
     assert stats.attacks_blocked == 1
 
 
+def test_on_detected_attack_wave(stats):
+    stats.on_detected_attack_wave(blocked=True)
+    assert stats.get_record()["requests"]["attackWaves"]["total"] == 1
+    assert stats.get_record()["requests"]["attackWaves"]["blocked"] == 1
+
+    stats.on_detected_attack_wave(blocked=False)
+    assert stats.get_record()["requests"]["attackWaves"]["total"] == 2
+    assert stats.get_record()["requests"]["attackWaves"]["blocked"] == 1
+
+
 def test_get_record(monkeypatch):
     # Mock the current time
     mock_time = 1234567890000
@@ -74,6 +89,8 @@ def test_get_record(monkeypatch):
     stats.on_detected_attack(blocked=True, operation="test.test")
     stats.attacks_detected = 5
     stats.attacks_blocked = 3
+    stats.on_detected_attack_wave(False)
+    stats.on_detected_attack_wave(False)
 
     record = stats.get_record()
     assert record["startedAt"] == stats.started_at
@@ -83,6 +100,8 @@ def test_get_record(monkeypatch):
     assert record["requests"]["aborted"] == 0
     assert record["requests"]["attacksDetected"]["total"] == 5
     assert record["requests"]["attacksDetected"]["blocked"] == 3
+    assert record["requests"]["attackWaves"]["total"] == 2
+    assert record["requests"]["attackWaves"]["blocked"] == 0
     assert record["operations"] == {
         "test.test": {
             "attacksDetected": {"blocked": 1, "total": 1},
@@ -153,6 +172,12 @@ def test_empty(stats):
     assert stats.empty() == False
 
 
+def test_empty_with_attack_waves(stats):
+    assert stats.empty()
+    stats.on_detected_attack_wave(blocked=False)
+    assert not stats.empty()
+
+
 def test_multiple_imports(stats):
     record1 = {
         "requests": {