From d5a26ca6eb119f76fcef1616eadee9452b8de015 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 29 Jun 2021 07:48:18 -0400 Subject: [PATCH] agent: Refactor BatchingTelemetryMessenger tests to destroy threads --- .../test_batching_telemetry_messenger.py | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py index 8b894a4f9..e3d2f89e0 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py @@ -1,7 +1,5 @@ import time -import pytest - from infection_monkey.telemetry.base_telem import BaseTelem from infection_monkey.telemetry.batchable_telem_mixin import BatchableTelemMixin from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem @@ -53,15 +51,21 @@ class BatchableTelemStub(BatchableTelemMixin, BaseTelem, IBatchableTelem): return {"entries": self._telemetry_entries} -@pytest.fixture -def batching_telemetry_messenger(monkeypatch, telemetry_messenger_spy): +# Note that this function is not a fixture. This is because BatchingTelemetyMessenger +# stops its thread when it is destructed. If this were a fixture, it may live +# past the end of the test, which would allow the in the BatchingTelemetryMessenger +# instance to keep running instead of stopping +def build_batching_telemetry_messenger(monkeypatch, telemetry_messenger_spy): patch_time(monkeypatch, 0) - return BatchingTelemetryMessenger(telemetry_messenger_spy, period=0.001) + return BatchingTelemetryMessenger(telemetry_messenger_spy, period=PERIOD) -def test_send_immediately(batching_telemetry_messenger, telemetry_messenger_spy): +def test_send_immediately(monkeypatch, telemetry_messenger_spy): + batching_telemetry_messenger = build_batching_telemetry_messenger( + monkeypatch, telemetry_messenger_spy + ) + telem = NonBatchableTelemStub() - batching_telemetry_messenger.send_telemetry(telem) release_GIL() @@ -69,7 +73,11 @@ def test_send_immediately(batching_telemetry_messenger, telemetry_messenger_spy) assert telemetry_messenger_spy.telemetries[0] == telem -def test_send_telem_batch(monkeypatch, batching_telemetry_messenger, telemetry_messenger_spy): +def test_send_telem_batch(monkeypatch, telemetry_messenger_spy): + batching_telemetry_messenger = build_batching_telemetry_messenger( + monkeypatch, telemetry_messenger_spy + ) + expected_data = {"entries": [1, 2]} telem1 = BatchableTelemStub(1) telem2 = BatchableTelemStub(2) @@ -86,9 +94,11 @@ def test_send_telem_batch(monkeypatch, batching_telemetry_messenger, telemetry_m assert telemetry_messenger_spy.telemetries[0].get_data() == expected_data -def test_send_different_telem_types( - monkeypatch, batching_telemetry_messenger, telemetry_messenger_spy -): +def test_send_different_telem_types(monkeypatch, telemetry_messenger_spy): + batching_telemetry_messenger = build_batching_telemetry_messenger( + monkeypatch, telemetry_messenger_spy + ) + telem1 = BatchableTelemStub(1, "cat1") telem2 = BatchableTelemStub(2, "cat2") @@ -105,7 +115,11 @@ def test_send_different_telem_types( assert telemetry_messenger_spy.telemetries[1] == telem2 -def test_send_two_batches(monkeypatch, batching_telemetry_messenger, telemetry_messenger_spy): +def test_send_two_batches(monkeypatch, telemetry_messenger_spy): + batching_telemetry_messenger = build_batching_telemetry_messenger( + monkeypatch, telemetry_messenger_spy + ) + telem1 = BatchableTelemStub(1, "cat1") telem2 = BatchableTelemStub(2, "cat1") @@ -125,9 +139,8 @@ def test_send_two_batches(monkeypatch, batching_telemetry_messenger, telemetry_m def test_send_remaining_telem_after_stop(monkeypatch, telemetry_messenger_spy): - patch_time(monkeypatch, 0) - batching_telemetry_messenger = BatchingTelemetryMessenger( - telemetry_messenger_spy, period=PERIOD + batching_telemetry_messenger = build_batching_telemetry_messenger( + monkeypatch, telemetry_messenger_spy ) expected_data = {"entries": [1]}