From 275efb2ab1531c34ed7ed98a1dde04ea7eda5f2f Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Thu, 22 Sep 2022 19:23:36 +0000
Subject: [PATCH] UT: Test on_terminate_agents_signal

---
 monkey/monkey_island/cc/services/__init__.py  |  2 +-
 .../cc/services/test_agent_signals_service.py | 27 ++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py
index c73aff356..8fdb5fc77 100644
--- a/monkey/monkey_island/cc/services/__init__.py
+++ b/monkey/monkey_island/cc/services/__init__.py
@@ -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
diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py
index 16aab943f..1beb53940 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py
+++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py
@@ -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)