From e48e2cb9af34077783755277c450599da6c9e1ad Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 11:38:52 -0400 Subject: [PATCH] UT: Add typehints to test_pba_file_upload.py --- .../cc/resources/test_pba_file_upload.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py index 423dd262c..a1f39aaf6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -1,9 +1,13 @@ +from typing import Callable + import pytest +from flask.testing import FlaskClient from tests.common import StubDIContainer from tests.monkey_island import InMemoryAgentConfigurationRepository, SingleFileRepository from tests.unit_tests.monkey_island.conftest import get_url_for_resource from tests.utils import raise_ +from common import DIContainer from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload @@ -25,17 +29,21 @@ Content-Type: text/x-python @pytest.fixture -def file_repository(): +def file_repository() -> IFileRepository: return SingleFileRepository() @pytest.fixture -def agent_configuration_repository(): +def agent_configuration_repository() -> IAgentConfigurationRepository: return InMemoryAgentConfigurationRepository() @pytest.fixture -def flask_client(build_flask_client, file_repository, agent_configuration_repository): +def flask_client( + build_flask_client: Callable[[DIContainer], FlaskClient], + file_repository: IFileRepository, + agent_configuration_repository: IAgentConfigurationRepository, +) -> FlaskClient: container = StubDIContainer() container.register_instance(IFileRepository, file_repository) container.register_instance(IAgentConfigurationRepository, agent_configuration_repository) @@ -45,7 +53,7 @@ def flask_client(build_flask_client, file_repository, agent_configuration_reposi @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_post(flask_client, pba_os): +def test_pba_file_upload_post(flask_client: FlaskClient, pba_os: str): url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( url, @@ -56,7 +64,7 @@ def test_pba_file_upload_post(flask_client, pba_os): assert resp.status_code == 200 -def test_pba_file_upload_post__invalid(flask_client): +def test_pba_file_upload_post__invalid(flask_client: FlaskClient): url = get_url_for_resource(FileUpload, target_os="bogus") resp = flask_client.post( url, @@ -68,7 +76,9 @@ def test_pba_file_upload_post__invalid(flask_client): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_post__internal_server_error(flask_client, pba_os, file_repository): +def test_pba_file_upload_post__internal_server_error( + flask_client: FlaskClient, pba_os: str, file_repository: IFileRepository +): file_repository.save_file = lambda x, y: raise_(Exception()) url = get_url_for_resource(FileUpload, target_os=pba_os) @@ -82,14 +92,14 @@ def test_pba_file_upload_post__internal_server_error(flask_client, pba_os, file_ @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_get__file_not_found(flask_client, pba_os): +def test_pba_file_upload_get__file_not_found(flask_client: FlaskClient, pba_os: str): url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") resp = flask_client.get(url) assert resp.status_code == 404 @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_file_download_endpoint_500(open_error_flask_client, pba_os): +def test_file_download_endpoint_500(open_error_flask_client, pba_os: str): url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") resp = open_error_flask_client.get(url) @@ -98,7 +108,7 @@ def test_file_download_endpoint_500(open_error_flask_client, pba_os): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_endpoint(flask_client, pba_os): +def test_pba_file_upload_endpoint(flask_client: FlaskClient, pba_os: str): url_with_os = get_url_for_resource(FileUpload, target_os=pba_os) resp_post = flask_client.post( @@ -123,7 +133,7 @@ def test_pba_file_upload_endpoint(flask_client, pba_os): assert resp_get_del.status_code == 404 -def test_pba_file_upload_endpoint__invalid(flask_client): +def test_pba_file_upload_endpoint__invalid(flask_client: FlaskClient): url_with_os = get_url_for_resource(FileUpload, target_os="bogus") resp_post = flask_client.post(