UT: Test on_terminate_agents_signal

This commit is contained in:
Kekoa Kaaikala 2022-09-22 19:23:36 +00:00 committed by Mike Salvatore
parent 9d3be7e1d3
commit 275efb2ab1
2 changed files with 27 additions and 2 deletions

View File

@ -1,4 +1,4 @@
from .agent_signals_service import AgentSignalsService
from .authentication_service import AuthenticationService
from .aws import AWSService
from .agent_signals_service import AgentSignalsService

View File

@ -4,7 +4,7 @@ from uuid import UUID
import pytest
from common.types import AgentID
from monkey_island.cc.models import Agent, Simulation
from monkey_island.cc.models import Agent, IslandMode, Simulation
from monkey_island.cc.repository import IAgentRepository, ISimulationRepository, UnknownRecordError
from monkey_island.cc.services import AgentSignalsService
@ -117,3 +117,28 @@ def test_progenitor_started_before_terminate(
signals = agent_signals_service.get_signals(agent.id)
assert signals.terminate.timestamp() == TERMINATE_TIMESTAMP
def test_on_terminate_agents_signal__stores_timestamp(
agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository
):
timestamp = 100
mock_simulation_repository.get_simulation = MagicMock(return_value=Simulation())
agent_signals_service.on_terminate_agents_signal(timestamp)
expected_value = Simulation(terminate_signal_time=timestamp)
assert mock_simulation_repository.save_simulation.called_once_with(expected_value)
def test_on_terminate_agents_signal__updates_timestamp(
agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository
):
timestamp = 100
mock_simulation_repository.get_simulation = MagicMock(
return_value=Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=50)
)
agent_signals_service.on_terminate_agents_signal(timestamp)
expected_value = Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=timestamp)
assert mock_simulation_repository.save_simulation.called_once_with(expected_value)