From b9576db426ef799d6a8782e2569dad5c77942cb8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 14:21:17 -0400 Subject: [PATCH] UT: Fix broken tests for find_available_island_apis() --- .../network/relay/test_utils.py | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py index 4979cce43..dd00ebc94 100644 --- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py @@ -2,7 +2,7 @@ import pytest import requests_mock from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIConnectionError -from infection_monkey.network.relay.utils import find_server +from infection_monkey.network.relay.utils import find_available_island_apis SERVER_1 = "1.1.1.1:12312" SERVER_2 = "2.2.2.2:4321" @@ -14,34 +14,42 @@ servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( - "expected_server, server_response_pairs", + "expected_available_servers, server_response_pairs", [ - (None, [(server, {"exc": IslandAPIConnectionError}) for server in servers]), + ([], [(server, {"exc": IslandAPIConnectionError}) for server in servers]), ( - SERVER_2, + servers[1:], [(SERVER_1, {"exc": IslandAPIConnectionError})] + [(server, {"text": ""}) for server in servers[1:]], # type: ignore[dict-item] ), ], ) -def test_find_server(expected_server, server_response_pairs): +def test_find_available_island_apis(expected_available_servers, server_response_pairs): with requests_mock.Mocker() as mock: for server, response in server_response_pairs: mock.get(f"https://{server}/api?action=is-up", **response) - actual_server, _ = find_server(servers) + available_apis = find_available_island_apis(servers) - assert actual_server is expected_server + assert len(available_apis) == len(server_response_pairs) + + for server, island_api_client in available_apis.items(): + if server in expected_available_servers: + assert island_api_client is not None + else: + assert island_api_client is None -def test_find_server__multiple_successes(): +def test_find_available_island_apis__multiple_successes(): + available_servers = [SERVER_2, SERVER_3] with requests_mock.Mocker() as mock: mock.get(f"https://{SERVER_1}/api?action=is-up", exc=IslandAPIConnectionError) - mock.get(f"https://{SERVER_2}/api?action=is-up", text="") - mock.get(f"https://{SERVER_3}/api?action=is-up", text="") - mock.get(f"https://{SERVER_4}/api?action=is-up", text="") + for server in available_servers: + mock.get(f"https://{server}/api?action=is-up", text="") - actual_server, actual_island_api_client = find_server(servers) + available_apis = find_available_island_apis(servers) - assert actual_server == SERVER_2 - assert isinstance(actual_island_api_client, IIslandAPIClient) + assert available_apis[SERVER_1] is None + assert available_apis[SERVER_4] is None + for server in available_servers: + assert isinstance(available_apis[server], IIslandAPIClient)