Merge pull request #2304 from guardicore/fix-slow-ransomware-test

Agent: Make BatchingTelemetryMessenger sleep period configurable

monkey/infection_monkey/payload/ransomware/ransomware_builder.py

@@ -1,6 +1,9 @@
 import logging
 from pprint import pformat
 
+from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
+    DEFAULT_PERIOD as DEFAULT_TELEMETRY_BATCH_PERIOD,
+)
 from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
     BatchingTelemetryMessenger,
 )
@@ -19,7 +22,12 @@ CHUNK_SIZE = 4096 * 24
 logger = logging.getLogger(__name__)
 
 
-def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger):
+def build_ransomware(
+    options: dict,
+    telemetry_messenger: ITelemetryMessenger,
+    # BatchingTelemetryMessenger will go away soon and so will this parameter
+    telemetry_batch_period: float = DEFAULT_TELEMETRY_BATCH_PERIOD,
+):
     logger.debug(f"Ransomware configuration:\n{pformat(options)}")
     ransomware_options = RansomwareOptions(options)
 
@@ -32,7 +40,7 @@ def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger):
         file_encryptor,
         file_selector,
         leave_readme,
-        BatchingTelemetryMessenger(telemetry_messenger),
+        BatchingTelemetryMessenger(telemetry_messenger, telemetry_batch_period),
     )
 
 

monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py

@@ -23,7 +23,9 @@ def test_uses_correct_extension(ransomware_options_dict, tmp_path, ransomware_fi
     ransomware_directories["linux_target_dir"] = target_dir
     ransomware_directories["windows_target_dir"] = target_dir
     telemetry_messenger = MagicMock()
-    ransomware = ransomware_builder.build_ransomware(ransomware_options_dict, telemetry_messenger)
+    ransomware = ransomware_builder.build_ransomware(
+        ransomware_options_dict, telemetry_messenger, 0.01
+    )
 
     file = target_dir / "file.txt"
     file.write_text("Do your worst!")