Island: Return island mode only as a string

monkey/monkey_island/cc/resources/island_mode.py

@@ -2,7 +2,7 @@ import json
 import logging
 from http import HTTPStatus
 
-from flask import make_response, request
+from flask import request
 
 from monkey_island.cc.models import IslandMode as IslandModeEnum
 from monkey_island.cc.resources.AbstractResource import AbstractResource
@@ -21,8 +21,7 @@ class IslandMode(AbstractResource):
     @jwt_required
     def put(self):
         try:
-            body = json.loads(request.data)
-            mode = IslandModeEnum(body.get("mode"))
+            mode = IslandModeEnum(json.loads(request.data))
 
             self._island_mode_service.set_mode(mode)
 
@@ -35,4 +34,4 @@ class IslandMode(AbstractResource):
     @jwt_required
     def get(self):
         island_mode = self._island_mode_service.get_mode()
-        return make_response({"mode": island_mode.value}, HTTPStatus.OK)
+        return island_mode.value, HTTPStatus.OK

monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py

@@ -38,24 +38,18 @@ def flask_client(build_flask_client):
 )
 def test_island_mode_post(flask_client, mode):
     resp = flask_client.put(
-        IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True
+        IslandModeResource.urls[0], data=json.dumps(mode), follow_redirects=True
     )
     assert resp.status_code == HTTPStatus.NO_CONTENT
 
 
 def test_island_mode_post__invalid_mode(flask_client):
     resp = flask_client.put(
-        IslandModeResource.urls[0], data=json.dumps({"mode": "bogus mode"}), follow_redirects=True
+        IslandModeResource.urls[0], data=json.dumps("bogus mode"), follow_redirects=True
     )
     assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
 
 
-@pytest.mark.parametrize("invalid_json", ["42", "{test"])
-def test_island_mode_post__invalid_json(flask_client, invalid_json):
-    resp = flask_client.put(IslandModeResource.urls[0], data="{test", follow_redirects=True)
-    assert resp.status_code == HTTPStatus.BAD_REQUEST
-
-
 def test_island_mode_post__internal_server_error(build_flask_client):
     mock_island_mode_service = MagicMock(spec=IslandModeService)
     mock_island_mode_service.set_mode = MagicMock(side_effect=RetrievalError)
@@ -66,7 +60,7 @@ def test_island_mode_post__internal_server_error(build_flask_client):
     with build_flask_client(container) as flask_client:
         resp = flask_client.put(
             IslandModeResource.urls[0],
-            data=json.dumps({"mode": IslandMode.RANSOMWARE.value}),
+            data=json.dumps(IslandMode.RANSOMWARE.value),
             follow_redirects=True,
         )
 
@@ -75,18 +69,16 @@ def test_island_mode_post__internal_server_error(build_flask_client):
 
 @pytest.mark.parametrize("mode", [IslandMode.RANSOMWARE.value, IslandMode.ADVANCED.value])
 def test_island_mode_endpoint(flask_client, mode):
-    flask_client.put(
-        IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True
-    )
+    flask_client.put(IslandModeResource.urls[0], data=json.dumps(mode), follow_redirects=True)
     resp = flask_client.get(IslandModeResource.urls[0], follow_redirects=True)
     assert resp.status_code == HTTPStatus.OK
-    assert json.loads(resp.data)["mode"] == mode
+    assert json.loads(resp.data) == mode
 
 
 def test_island_mode_endpoint__invalid_mode(flask_client):
     resp_post = flask_client.put(
-        IslandModeResource.urls[0], data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True
+        IslandModeResource.urls[0], data=json.dumps("bogus_mode"), follow_redirects=True
     )
     resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True)
     assert resp_post.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
-    assert json.loads(resp_get.data)["mode"] == IslandMode.UNSET.value
+    assert json.loads(resp_get.data) == IslandMode.UNSET.value