diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index bd6bfcb41..b18f3990b 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -13,10 +13,12 @@ from infection_monkey.island_api_client import ( SERVER = "1.1.1.1:9999" PBA_FILE = "dummy.pba" +WINDOWS = "windows" ISLAND_URI = f"https://{SERVER}/api?action=is-up" ISLAND_SEND_LOG_URI = f"https://{SERVER}/api/log" ISLAND_GET_PBA_FILE_URI = f"https://{SERVER}/api/pba/download/{PBA_FILE}" +ISLAND_GET_AGENT_BINARY_URI = f"https://{SERVER}/api/agent-binaries/{WINDOWS}" @pytest.mark.parametrize( @@ -118,3 +120,38 @@ def test_island_api_client_get_pba_file__status_code(status_code, expected_error with pytest.raises(expected_error): m.get(ISLAND_GET_PBA_FILE_URI, status_code=status_code) island_api_client.get_pba_file(filename=PBA_FILE) + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + (Exception, IslandAPIError), + ], +) +def test_island_api_client__get_agent_binary(actual_error, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_AGENT_BINARY_URI, exc=actual_error) + island_api_client.get_agent_binary(os_name=WINDOWS) + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client__get_agent_binary_status_code(status_code, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_AGENT_BINARY_URI, status_code=status_code) + island_api_client.get_agent_binary(os_name=WINDOWS)