Island: Add unset mode to get method in island mode

This commit is contained in:
Ilija Lazoroski 2021-07-14 12:40:08 +02:00
parent bf517bf566
commit 2a1d41f6c7
3 changed files with 20 additions and 9 deletions

View File

@ -5,7 +5,7 @@ import flask_restful
from flask import make_response, request from flask import make_response, request
from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.resources.auth.auth import jwt_required
from monkey_island.cc.services.mode.island_mode_service import get_mode, set_mode from monkey_island.cc.services.mode.island_mode_service import ModeNotSetError, get_mode, set_mode
from monkey_island.cc.services.mode.mode_enum import IslandModeEnum from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -32,5 +32,6 @@ class IslandMode(flask_restful.Resource):
try: try:
island_mode = get_mode() island_mode = get_mode()
return make_response({"mode": island_mode}, 200) return make_response({"mode": island_mode}, 200)
except IndexError:
return make_response({}, 422) except ModeNotSetError:
return make_response({"mode": None}, 200)

View File

@ -8,6 +8,15 @@ def set_mode(mode: IslandModeEnum):
island_mode_model.save() island_mode_model.save()
def get_mode(): def get_mode() -> str:
mode = IslandMode.objects[0].mode if IslandMode.objects:
return mode mode = IslandMode.objects[0].mode
return mode
else:
raise ModeNotSetError
class ModeNotSetError(Exception):
"""
Throw this exception when island mode is not set.
"""

View File

@ -51,8 +51,9 @@ def test_island_mode_endpoint(flask_client, uses_database, mode):
def test_island_mode_endpoint__invalid_mode(flask_client, uses_database): def test_island_mode_endpoint__invalid_mode(flask_client, uses_database):
flask_client.post( resp_post = flask_client.post(
"/api/island-mode", data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True "/api/island-mode", data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True
) )
resp = flask_client.get("/api/island-mode", follow_redirects=True) resp_get = flask_client.get("/api/island-mode", follow_redirects=True)
assert resp.status_code == 422 assert resp_post.status_code == 422
assert json.loads(resp_get.data)["mode"] is None