UT: add unit tests for monkey killing and monkey parent fetching

monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py

@@ -3,7 +3,7 @@ import uuid
 
 import pytest
 
-from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError
+from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError, ParentNotFoundError
 from monkey_island.cc.models.monkey_ttl import MonkeyTtl
 
 logger = logging.getLogger(__name__)
@@ -162,3 +162,35 @@ class TestMonkey:
 
         cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info()
         assert cache_info_after_query.hits == 2
+
+    @pytest.mark.usefixtures("uses_database")
+    def test_has_parent(self):
+        monkey_1 = Monkey(guid=str(uuid.uuid4()))
+        monkey_2 = Monkey(guid=str(uuid.uuid4()))
+        monkey_1.parent = [[monkey_2.guid]]
+        monkey_1.save()
+        assert monkey_1.has_parent()
+
+    @pytest.mark.usefixtures("uses_database")
+    def test_has_no_parent(self):
+        monkey_1 = Monkey(guid=str(uuid.uuid4()))
+        monkey_1.parent = [[monkey_1.guid]]
+        monkey_1.save()
+        assert not monkey_1.has_parent()
+
+    @pytest.mark.usefixtures("uses_database")
+    def test_get_parent(self):
+        monkey_1 = Monkey(guid=str(uuid.uuid4()))
+        monkey_2 = Monkey(guid=str(uuid.uuid4()))
+        monkey_1.parent = [[monkey_2.guid]]
+        monkey_1.save()
+        monkey_2.save()
+        assert monkey_1.get_parent().guid == monkey_2.guid
+
+    @pytest.mark.usefixtures("uses_database")
+    def test_get_parent_no_parent(self):
+        monkey_1 = Monkey(guid=str(uuid.uuid4()))
+        monkey_1.parent = [[monkey_1.guid]]
+        monkey_1.save()
+        with pytest.raises(ParentNotFoundError):
+            monkey_1.get_parent()

monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py

@@ -0,0 +1,101 @@
+import uuid
+
+import pytest
+
+from monkey_island.cc.models import Config, Monkey
+from monkey_island.cc.models.agent_controls import AgentControls
+from monkey_island.cc.services.infection_lifecycle import should_agent_die
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_should_agent_die_by_config(monkeypatch):
+    monkey = Monkey(guid=str(uuid.uuid4()))
+    monkey.config = Config(alive=False)
+    monkey.save()
+    assert should_agent_die(monkey.guid)
+
+    monkeypatch.setattr(
+        "monkey_island.cc.services.infection_lifecycle._is_monkey_killed_manually", lambda _: False
+    )
+    monkey.config.alive = True
+    monkey.save()
+    assert not should_agent_die(monkey.guid)
+
+
+def create_monkey(launch_time):
+    monkey = Monkey(guid=str(uuid.uuid4()))
+    monkey.config = Config(alive=True)
+    monkey.launch_time = launch_time
+    monkey.save()
+    return monkey
+
+
+def create_kill_event(event_time):
+    kill_event = AgentControls(last_stop_all=event_time)
+    kill_event.save()
+    return kill_event
+
+
+def create_parent(child_monkey, launch_time):
+    monkey_parent = Monkey(guid=str(uuid.uuid4()))
+    child_monkey.parent = [[monkey_parent.guid]]
+    monkey_parent.launch_time = launch_time
+    monkey_parent.save()
+    child_monkey.save()
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_was_agent_killed_manually(monkeypatch):
+    monkey = create_monkey(launch_time=2)
+
+    create_kill_event(event_time=3)
+
+    assert should_agent_die(monkey.guid)
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_agent_killed_on_wakeup(monkeypatch):
+    monkey = create_monkey(launch_time=2)
+
+    create_kill_event(event_time=2)
+
+    assert should_agent_die(monkey.guid)
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_manual_kill_dont_affect_new_monkeys(monkeypatch):
+    monkey = create_monkey(launch_time=3)
+
+    create_kill_event(event_time=2)
+
+    assert not should_agent_die(monkey.guid)
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_parent_manually_killed(monkeypatch):
+    monkey = create_monkey(launch_time=3)
+    create_parent(child_monkey=monkey, launch_time=1)
+
+    create_kill_event(event_time=2)
+
+    assert should_agent_die(monkey.guid)
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_parent_manually_killed_on_wakeup(monkeypatch):
+    monkey = create_monkey(launch_time=3)
+    create_parent(child_monkey=monkey, launch_time=2)
+
+    create_kill_event(event_time=2)
+
+    assert should_agent_die(monkey.guid)
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_manual_kill_dont_affect_new_monkeys_with_parent(monkeypatch):
+    monkey = create_monkey(launch_time=3)
+    create_parent(child_monkey=monkey, launch_time=2)
+
+    create_kill_event(event_time=1)
+
+    assert not should_agent_die(monkey.guid)