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 http import HTTPStatus
from unittest.mock import MagicMock from unittest.mock import MagicMock
@ -39,8 +38,7 @@ def flask_client(build_flask_client):
def test_island_mode_post(flask_client, mode): def test_island_mode_post(flask_client, mode):
resp = flask_client.put( resp = flask_client.put(
IslandModeResource.urls[0], IslandModeResource.urls[0],
data=json.dumps(mode), json=mode,
headers={"Content-Type": "application/json"},
follow_redirects=True, follow_redirects=True,
) )
assert resp.status_code == HTTPStatus.NO_CONTENT 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): def test_island_mode_post__invalid_mode(flask_client):
resp = flask_client.put( resp = flask_client.put(
IslandModeResource.urls[0], IslandModeResource.urls[0],
data=json.dumps("bogus mode"), json="bogus mode",
headers={"Content-Type": "application/json"},
follow_redirects=True, follow_redirects=True,
) )
assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY 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: with build_flask_client(container) as flask_client:
resp = flask_client.put( resp = flask_client.put(
IslandModeResource.urls[0], IslandModeResource.urls[0],
data=json.dumps(IslandMode.RANSOMWARE.value), json=IslandMode.RANSOMWARE.value,
headers={"Content-Type": "application/json"},
follow_redirects=True, 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): def test_island_mode_endpoint(flask_client, mode):
flask_client.put( flask_client.put(
IslandModeResource.urls[0], IslandModeResource.urls[0],
data=json.dumps(mode), json=mode,
headers={"Content-Type": "application/json"},
follow_redirects=True, follow_redirects=True,
) )
resp = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) resp = flask_client.get(IslandModeResource.urls[0], follow_redirects=True)
assert resp.status_code == HTTPStatus.OK 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): def test_island_mode_endpoint__invalid_mode(flask_client):
resp_post = flask_client.put( resp_post = flask_client.put(
IslandModeResource.urls[0], IslandModeResource.urls[0],
data=json.dumps("bogus_mode"), json="bogus_mode",
headers={"Content-Type": "application/json"},
follow_redirects=True, follow_redirects=True,
) )
resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True)
assert resp_post.status_code == HTTPStatus.UNPROCESSABLE_ENTITY assert resp_post.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
assert json.loads(resp_get.data) == IslandMode.UNSET.value assert resp_get.json == IslandMode.UNSET.value