Island: Add UT tests for island mode model

This commit is contained in:
Ilija Lazoroski 2021-07-13 10:30:12 +02:00
parent 3bde6f013a
commit f9ed53a527
5 changed files with 58 additions and 10 deletions

View File

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

View File

@ -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")
try:
mode = IslandModeEnum(mode_str)
island_mode_service.set_mode(mode)
# TODO return status
return make_response({})
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)

View File

@ -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

View File

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

View File

@ -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