Agent: Make BatchingTelemetryMessenger sleep period configurable
The ransomware integration tests were sometimes painfully slow because they were waiting for the BatchingTelemetryMessenger to stop its internal thread. This commit surfaces the sleep period parameter to the `build_ransomware()` function so that the unit tests can set it to a very small value. This is somewhat of a hack, but BatchingTelemetryMessenger is going away very soon, as Telemetry is being replaced by Events.
This commit is contained in:
parent
f6270400b6
commit
f677096439
|
@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
@ -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!")
|
||||
|
|
Loading…
Reference in New Issue