UT: Simplify some logic in test_island_mode.py

This commit is contained in:
Shreya Malviya 2022-08-03 18:14:46 +05:30
parent 4bcab891b2
commit 5a06e88913
1 changed files with 7 additions and 13 deletions

View File

@ -1,4 +1,3 @@
import json
from http import HTTPStatus
from unittest.mock import MagicMock
@ -39,8 +38,7 @@ 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),
headers={"Content-Type": "application/json"},
json=mode,
follow_redirects=True,
)
assert resp.status_code == HTTPStatus.NO_CONTENT
@ -49,8 +47,7 @@ def test_island_mode_post(flask_client, mode):
def test_island_mode_post__invalid_mode(flask_client):
resp = flask_client.put(
IslandModeResource.urls[0],
data=json.dumps("bogus mode"),
headers={"Content-Type": "application/json"},
json="bogus mode",
follow_redirects=True,
)
assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
@ -66,8 +63,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(IslandMode.RANSOMWARE.value),
headers={"Content-Type": "application/json"},
json=IslandMode.RANSOMWARE.value,
follow_redirects=True,
)
@ -78,22 +74,20 @@ def test_island_mode_post__internal_server_error(build_flask_client):
def test_island_mode_endpoint(flask_client, mode):
flask_client.put(
IslandModeResource.urls[0],
data=json.dumps(mode),
headers={"Content-Type": "application/json"},
json=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
assert resp.json == mode
def test_island_mode_endpoint__invalid_mode(flask_client):
resp_post = flask_client.put(
IslandModeResource.urls[0],
data=json.dumps("bogus_mode"),
headers={"Content-Type": "application/json"},
json="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) == IslandMode.UNSET.value
assert resp_get.json == IslandMode.UNSET.value