From f9ed53a5272de4e5c50b5d0423b411288c993285 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Tue, 13 Jul 2021 10:30:12 +0200
Subject: [PATCH] Island: Add UT tests for island mode model

---
 .../cc/models/island_mode_model.py            |  5 ++++
 .../monkey_island/cc/resources/island_mode.py | 16 +++++++----
 .../cc/services/mode/island_mode_service.py   |  6 +++-
 .../monkey_island/cc/resources/conftest.py    | 13 +++++++++
 .../cc/resources/test_island_mode.py          | 28 +++++++++++++++++--
 5 files changed, 58 insertions(+), 10 deletions(-)
 create mode 100644 monkey/monkey_island/cc/models/island_mode_model.py

diff --git a/monkey/monkey_island/cc/models/island_mode_model.py b/monkey/monkey_island/cc/models/island_mode_model.py
new file mode 100644
index 000000000..dec93e501
--- /dev/null
+++ b/monkey/monkey_island/cc/models/island_mode_model.py
@@ -0,0 +1,5 @@
+from mongoengine import Document, StringField
+
+
+class IslandMode(Document):
+    mode = StringField()
diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py
index d0a109564..cf8a3679e 100644
--- a/monkey/monkey_island/cc/resources/island_mode.py
+++ b/monkey/monkey_island/cc/resources/island_mode.py
@@ -1,20 +1,24 @@
 import json
+import logging
 
 import flask_restful
 from flask import make_response, request
 
 from monkey_island.cc.resources.auth.auth import jwt_required
-from monkey_island.cc.services.mode import island_mode_service
+from monkey_island.cc.services.mode.island_mode_service import set_mode
 from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
 
+logger = logging.getLogger(__name__)
+
 
 class IslandMode(flask_restful.Resource):
     @jwt_required
     def post(self):
         body = json.loads(request.data)
         mode_str = body.get("mode")
-        mode = IslandModeEnum(mode_str)
-        island_mode_service.set_mode(mode)
-
-        # TODO return status
-        return make_response({})
+        try:
+            mode = IslandModeEnum(mode_str)
+            mode_value = set_mode(mode)
+            return make_response({"status": "MODE_FOUND", "mode": mode_value}, 200)
+        except ValueError:
+            return make_response({"status": "MODE_NOT_FOUND"}, 404)
diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py
index 86f3e6d09..75041d2c1 100644
--- a/monkey/monkey_island/cc/services/mode/island_mode_service.py
+++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py
@@ -1,5 +1,9 @@
+from monkey_island.cc.models.island_mode_model import IslandMode
 from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
 
 
 def set_mode(mode: IslandModeEnum):
-    pass
+    island_mode_model = IslandMode()
+    island_mode_model.mode = mode.value
+    island_mode_model.save()
+    return mode.value
diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py
index 0e82fe163..2421e00aa 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py
+++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py
@@ -1,11 +1,13 @@
 import flask_jwt_extended
 import flask_restful
+import mongoengine
 import pytest
 from flask import Flask
 
 import monkey_island.cc.app
 import monkey_island.cc.resources.auth.auth
 import monkey_island.cc.resources.island_mode
+from monkey_island.cc.models.island_mode_model import IslandMode
 from monkey_island.cc.services.representations import output_json
 
 
@@ -30,3 +32,14 @@ def mock_init_app():
     monkey_island.cc.app.init_api_resources(api)
 
     return app
+
+
+@pytest.fixture(scope="module", autouse=True)
+def fake_mongo():
+    mongoengine.disconnect()
+    mongoengine.connect("mongoenginetest", host="mongomock://localhost")
+
+
+@pytest.fixture(scope="function")
+def uses_database():
+    IslandMode.objects().delete()
diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py
index 01cb17aea..13ded4799 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py
+++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py
@@ -2,12 +2,34 @@ import json
 
 import pytest
 
+from monkey_island.cc.models.island_mode_model import IslandMode
+
 
 def test_island_mode_post(flask_client):
-    resp = flask_client.post('/api/island-mode', data=json.dumps({"mode": "ransomware"}), follow_redirects=True)
+    resp = flask_client.post(
+        "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True
+    )
     assert resp.status_code == 200
 
 
 def test_island_mode_post__invalid_mode(flask_client):
-    with pytest.raises(TypeError):
-        flask_client.post('/api/island-mode', data=json.dumps({"mode": "bogus mode"}), follow_redirects=True)
+    resp = flask_client.post(
+        "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True
+    )
+    assert resp.status_code == 404
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_island_mode_post__set_model(flask_client):
+    flask_client.post(
+        "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True
+    )
+    assert IslandMode.objects[0].mode == "ransomware"
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_island_mode_post__set_invalid_model(flask_client):
+    flask_client.post(
+        "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True
+    )
+    assert len(IslandMode.objects) == 0