UT: Add helper method to form URL's

This method forms url's based on parameters and allows us to avoid hard-coding url's in our unit tests
This commit is contained in:
vakarisz 2022-05-23 14:40:45 +03:00
parent 75318059e4
commit 4069cc8084
3 changed files with 58 additions and 17 deletions

View File

@ -3,7 +3,9 @@ from typing import BinaryIO
import pytest import pytest
from tests.common import StubDIContainer from tests.common import StubDIContainer
from tests.unit_tests.monkey_island.conftest import get_url_for_resource
from monkey_island.cc.resources.pba_file_download import PBAFileDownload
from monkey_island.cc.services import FileRetrievalError, IFileStorageService from monkey_island.cc.services import FileRetrievalError, IFileStorageService
FILE_NAME = "test_file" FILE_NAME = "test_file"
@ -40,7 +42,8 @@ def flask_client(build_flask_client):
def test_file_download_endpoint(tmp_path, flask_client): def test_file_download_endpoint(tmp_path, flask_client):
resp = flask_client.get(f"/api/pba/download/{FILE_NAME}") download_url = get_url_for_resource(PBAFileDownload, filename=FILE_NAME)
resp = flask_client.get(download_url)
assert resp.status_code == 200 assert resp.status_code == 200
assert next(resp.response) == FILE_CONTENTS assert next(resp.response) == FILE_CONTENTS
@ -48,7 +51,8 @@ def test_file_download_endpoint(tmp_path, flask_client):
def test_file_download_endpoint_404(tmp_path, flask_client): def test_file_download_endpoint_404(tmp_path, flask_client):
nonexistant_file_name = "nonexistant_file" nonexistant_file_name = "nonexistant_file"
download_url = get_url_for_resource(PBAFileDownload, filename=nonexistant_file_name)
resp = flask_client.get(f"/api/pba/download/{nonexistant_file_name}") resp = flask_client.get(download_url)
assert resp.status_code == 404 assert resp.status_code == 404

View File

@ -3,9 +3,10 @@ from typing import BinaryIO
import pytest import pytest
from tests.common import StubDIContainer from tests.common import StubDIContainer
from tests.unit_tests.monkey_island.conftest import get_url_for_resource
from tests.utils import raise_ from tests.utils import raise_
from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload
from monkey_island.cc.services import FileRetrievalError, IFileStorageService from monkey_island.cc.services import FileRetrievalError, IFileStorageService
TEST_FILE_CONTENTS = b"m0nk3y" TEST_FILE_CONTENTS = b"m0nk3y"
@ -74,8 +75,9 @@ def flask_client(build_flask_client, file_storage_service):
@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE])
def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value): def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value):
url = get_url_for_resource(FileUpload, target_os=pba_os)
resp = flask_client.post( resp = flask_client.post(
f"/api/file-upload/{pba_os}", url,
data=TEST_FILE, data=TEST_FILE,
content_type="multipart/form-data; " "boundary=---------------------------" "1", content_type="multipart/form-data; " "boundary=---------------------------" "1",
follow_redirects=True, follow_redirects=True,
@ -84,8 +86,9 @@ def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value):
def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value): def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value):
url = get_url_for_resource(FileUpload, target_os="bogus")
resp = flask_client.post( resp = flask_client.post(
"/api/file-upload/bogus", url,
data=TEST_FILE, data=TEST_FILE,
content_type="multipart/form-data; " "boundary=---------------------------" "1", content_type="multipart/form-data; " "boundary=---------------------------" "1",
follow_redirects=True, follow_redirects=True,
@ -98,9 +101,10 @@ def test_pba_file_upload_post__internal_server_error(
flask_client, pba_os, mock_set_config_value, file_storage_service flask_client, pba_os, mock_set_config_value, file_storage_service
): ):
file_storage_service.save_file = lambda x, y: raise_(Exception()) file_storage_service.save_file = lambda x, y: raise_(Exception())
url = get_url_for_resource(FileUpload, target_os=pba_os)
resp = flask_client.post( resp = flask_client.post(
f"/api/file-upload/{pba_os}", url,
data=TEST_FILE, data=TEST_FILE,
content_type="multipart/form-data; boundary=---------------------------1", content_type="multipart/form-data; boundary=---------------------------1",
follow_redirects=True, follow_redirects=True,
@ -110,7 +114,8 @@ def test_pba_file_upload_post__internal_server_error(
@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE])
def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_config_value): def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_config_value):
resp = flask_client.get(f"/api/file-upload/{pba_os}?load=bogus_mogus.py") url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py")
resp = flask_client.get(url)
assert resp.status_code == 404 assert resp.status_code == 404
@ -118,23 +123,24 @@ def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_conf
def test_pba_file_upload_endpoint( def test_pba_file_upload_endpoint(
flask_client, pba_os, mock_get_config_value, mock_set_config_value flask_client, pba_os, mock_get_config_value, mock_set_config_value
): ):
url_with_os = get_url_for_resource(FileUpload, target_os=pba_os)
resp_post = flask_client.post( resp_post = flask_client.post(
f"/api/file-upload/{pba_os}", url_with_os,
data=TEST_FILE, data=TEST_FILE,
content_type="multipart/form-data; " "boundary=---------------------------" "1", content_type="multipart/form-data; " "boundary=---------------------------" "1",
follow_redirects=True, follow_redirects=True,
) )
resp_get = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py") url_with_filename = get_url_for_resource(FileUpload, target_os=pba_os, filename="test.py")
resp_get = flask_client.get(url_with_filename)
assert resp_get.status_code == 200 assert resp_get.status_code == 200
assert resp_get.data == TEST_FILE_CONTENTS assert resp_get.data == TEST_FILE_CONTENTS
# Closing the response closes the file handle, else it can't be deleted # Closing the response closes the file handle, else it can't be deleted
resp_get.close() resp_get.close()
resp_delete = flask_client.delete( resp_delete = flask_client.delete(url_with_os, data="test.py", content_type="text/plain;")
f"/api/file-upload/{pba_os}", data="test.py", content_type="text/plain;" resp_get_del = flask_client.get(url_with_filename)
)
resp_get_del = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py")
assert resp_post.status_code == 200 assert resp_post.status_code == 200
assert resp_delete.status_code == 200 assert resp_delete.status_code == 200
@ -145,16 +151,20 @@ def test_pba_file_upload_endpoint(
def test_pba_file_upload_endpoint__invalid( def test_pba_file_upload_endpoint__invalid(
flask_client, mock_set_config_value, mock_get_config_value flask_client, mock_set_config_value, mock_get_config_value
): ):
url_with_os = get_url_for_resource(FileUpload, target_os="bogus")
resp_post = flask_client.post( resp_post = flask_client.post(
"/api/file-upload/bogus", url_with_os,
data=TEST_FILE, data=TEST_FILE,
content_type="multipart/form-data; " "boundary=---------------------------" "1", content_type="multipart/form-data; " "boundary=---------------------------" "1",
follow_redirects=True, follow_redirects=True,
) )
resp_get = flask_client.get("/api/file-upload/bogus?load=test.py")
resp_delete = flask_client.delete( url_with_filename = get_url_for_resource(
"/api/file-upload/bogus", data="test.py", content_type="text/plain;" FileUpload, target_os="bogus", filename="bobug_mogus.py"
) )
resp_get = flask_client.get(url_with_filename)
resp_delete = flask_client.delete(url_with_os, data="test.py", content_type="text/plain;")
assert resp_post.status_code == 422 assert resp_post.status_code == 422
assert resp_get.status_code == 422 assert resp_get.status_code == 422
assert resp_delete.status_code == 422 assert resp_delete.status_code == 422

View File

@ -1,11 +1,14 @@
import os import os
import re
from collections.abc import Callable from collections.abc import Callable
from typing import Set
import flask_restful import flask_restful
import pytest import pytest
from flask import Flask from flask import Flask
import monkey_island import monkey_island
from monkey_island.cc.resources.i_resource import IResource
from monkey_island.cc.services.representations import output_json from monkey_island.cc.services.representations import output_json
@ -37,3 +40,27 @@ def mock_flask_resource_manager(container):
flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container) flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container)
return flask_resource_manager return flask_resource_manager
def get_url_for_resource(resource: IResource, **kwargs):
chosen_url = None
for url in resource.urls:
if _get_url_keywords(url) == set(kwargs.keys()):
chosen_url = url
if not chosen_url:
raise Exception(
f"Resource {resource} doesn't contain a url that matches {kwargs} keywords."
)
for key, value in kwargs.items():
reg_pattern = f"<.*:{key}>"
chosen_url = re.sub(pattern=reg_pattern, repl=value, string=chosen_url)
return chosen_url
def _get_url_keywords(url: str) -> Set[str]:
# Match pattern <something:keyword>, but only put "keyword" in a group
reg_pattern = "(?:<.*?:)(.*?)(?:>)"
reg_matches = re.finditer(reg_pattern, url)
return set([match.groups()[0] for match in reg_matches])