Fix test cases for new put_event data structure changes

Author: bitterpanda63

aikido_zen/background_process/cloud_connection_manager/on_start_test.py

@@ -8,7 +8,7 @@ def mock_connection_manager():
     connection_manager = MagicMock()
     connection_manager.token = "test_token"
     connection_manager.timeout_in_sec = 5
-    connection_manager.api.report = MagicMock(return_value={"success": True})
+    connection_manager.report_api_event = MagicMock(return_value={"success": True})
     connection_manager.get_manager_info = lambda: {}
     connection_manager.update_service_config = MagicMock()
     return connection_manager
@@ -19,15 +19,15 @@ def test_on_start_no_token():
     connection_manager = MagicMock()
     connection_manager.token = None
     on_start(connection_manager)
-    connection_manager.api.report.assert_not_called()
+    connection_manager.report_api_event.assert_called()
 
 
 def test_on_start_success(mock_connection_manager, caplog):
     """Test that the API call is made successfully and the service config is updated."""
     on_start(mock_connection_manager)
 
     # Check that the API report method was called
-    mock_connection_manager.api.report.assert_called_once()
+    mock_connection_manager.report_api_event.assert_called_once()
 
     # Check that the service config was updated
     mock_connection_manager.update_service_config.assert_called_once()
@@ -38,18 +38,15 @@ def test_on_start_success(mock_connection_manager, caplog):
 
 def test_on_start_failure(mock_connection_manager, caplog):
     """Test that an error is logged when the API call fails."""
-    mock_connection_manager.api.report.return_value = {
+    mock_connection_manager.report_api_event.return_value = {
         "success": False,
         "error": "Some error",
     }
 
     on_start(mock_connection_manager)
 
     # Check that the API report method was called
-    mock_connection_manager.api.report.assert_called_once()
+    mock_connection_manager.report_api_event.assert_called_once()
 
     # Check that the service config was not updated
     mock_connection_manager.update_service_config.assert_not_called()
-
-    # Check that the error log was called
-    assert "Failed to communicate with Aikido Server : Some error" in caplog.text

aikido_zen/background_process/commands/sync_data_test.py

@@ -72,7 +72,7 @@ def test_process_sync_data_initialization(setup_connection_manager):
         ],
     }
 
-    result = process_sync_data(connection_manager, data, None)
+    result = process_sync_data(connection_manager, data)
 
     # Check that routes were initialized correctly
     assert len(connection_manager.routes) == 2
@@ -147,7 +147,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
         "middleware_installed": True,
     }
 
-    result = process_sync_data(connection_manager, data, None)
+    result = process_sync_data(connection_manager, data)
 
     # Check that routes were initialized correctly
     assert len(connection_manager.routes) == 2
@@ -214,7 +214,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
     }
 
     # First call to initialize the route
-    process_sync_data(connection_manager, data, None)
+    process_sync_data(connection_manager, data)
 
     # Second call to update the existing route
     data_update = {
@@ -241,7 +241,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
         },
     }
 
-    result = process_sync_data(connection_manager, data_update, None)
+    result = process_sync_data(connection_manager, data_update)
 
     # Check that the hit count was updated correctly
     assert (
@@ -275,7 +275,7 @@ def test_process_sync_data_no_routes(setup_connection_manager):
     connection_manager = setup_connection_manager
     data = {"current_routes": {}, "reqs": 0}  # No requests to add
 
-    result = process_sync_data(connection_manager, data, None)
+    result = process_sync_data(connection_manager, data)
 
     # Check that no routes were initialized
     assert len(connection_manager.routes) == 0

aikido_zen/sources/functions/request_handler_test.py

@@ -42,7 +42,7 @@ def __init__(self):
     def send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1):
         if action != "CHECK_FIREWALL_LISTS":
             return {"success": False}
-        res = process_check_firewall_lists(self.conn_manager, obj, None, None)
+        res = process_check_firewall_lists(self.conn_manager, obj)
         return {
             "success": True,
             "data": res,

aikido_zen/vulnerabilities/init_test.py

@@ -144,13 +144,29 @@ def test_sql_injection_with_comms(caplog, get_context, monkeypatch):
             )
         mock_comms.send_data_to_bg_process.assert_called_once()
         call_args = mock_comms.send_data_to_bg_process.call_args[0]
-        assert call_args[0] == "ATTACK"
-        assert call_args[1][0]["kind"] == "sql_injection"
-        assert (
-            call_args[1][0]["metadata"]["sql"]
-            == "INSERT * INTO VALUES ('doggoss2', TRUE);"
-        )
-        assert call_args[1][0]["metadata"]["dialect"] == "mysql"
+        assert call_args[0] == "put_event"
+        assert call_args[1].event["request"] == {
+            "ipAddress": "198.51.100.23",
+            "method": "GET",
+            "route": "/hello",
+            "source": "flask",
+            "url": "http://localhost:8080/hello",
+            "userAgent": None,
+        }
+        del call_args[1].event["attack"]["stack"]  # Hard to test
+        assert call_args[1].event["attack"] == {
+            "blocked": True,
+            "kind": "sql_injection",
+            "metadata": {
+                "dialect": "mysql",
+                "sql": "INSERT * INTO VALUES ('doggoss2', TRUE);",
+            },
+            "operation": "test_op",
+            "pathToPayload": ".test_input_sql",
+            "payload": '"doggoss2\', TRUE"',
+            "source": "body",
+            "user": None,
+        }
 
 
 def test_ssrf_vulnerability_scan_adds_hostname(get_context):