Island: Add UT tests for island mode model

monkey/monkey_island/cc/models/island_mode_model.py

@@ -0,0 +1,5 @@
+from mongoengine import Document, StringField
+
+
+class IslandMode(Document):
+    mode = StringField()

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)

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

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()

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