p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

UT: Add proper test for ControlChannel.should_agent_stop()

This commit is contained in:
Mike Salvatore 2022-09-23 12:25:09 -04:00
parent a5f1117ce3
commit f7198ea98a
1 changed files with 14 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py
View File

@ -1,8 +1,10 @@
from typing import Optional
from unittest.mock import MagicMock from unittest.mock import MagicMock
import pytest import pytest
from infection_monkey.i_control_channel import IslandCommunicationError from common import AgentSignals
from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
from infection_monkey.island_api_client import ( from infection_monkey.island_api_client import (
IIslandAPIClient, IIslandAPIClient,
IslandAPIConnectionError, IslandAPIConnectionError,
@ -33,9 +35,17 @@ def control_channel(island_api_client) -> ControlChannel:
return ControlChannel(SERVER, AGENT_ID, island_api_client) return ControlChannel(SERVER, AGENT_ID, island_api_client)
def test_control_channel__should_agent_stop(control_channel, island_api_client): @pytest.mark.parametrize("signal_time,expected_should_stop", [(1663950115, True), (None, False)])
control_channel.should_agent_stop() def test_control_channel__should_agent_stop(
assert island_api_client.get_agent_signals.called_once() control_channel: IControlChannel,
island_api_client: IIslandAPIClient,
signal_time: Optional[int],
expected_should_stop: bool,
):
island_api_client.get_agent_signals = MagicMock(
return_value=AgentSignals(terminate=signal_time)
)
assert control_channel.should_agent_stop() is expected_should_stop
@pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) @pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS)