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 a013fcec7..27d31610b 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 @@ -41,6 +41,7 @@ ISLAND_SEND_EVENTS_URI = f"https://{SERVER}/api/agent-events" ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents" ISLAND_AGENT_STOP_URI = f"https://{SERVER}/api/monkey-control/needs-to-stop/{AGENT_ID}" ISLAND_GET_CONFIG_URI = f"https://{SERVER}/api/agent-configuration" +ISLAND_GET_PROPAGATION_CREDENTIALS_URI = f"https://{SERVER}/api/propagation-credentials" class Event1(AbstractAgentEvent): @@ -408,3 +409,49 @@ def test_island_api_client_get_config__bad_json(): with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_GET_CONFIG_URI, content=b"bad") island_api_client.get_config() + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__get_credentials_for_propagation(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_PROPAGATION_CREDENTIALS_URI, exc=actual_error) + island_api_client.get_credentials_for_propagation() + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_get_credentials_for_propagation__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_PROPAGATION_CREDENTIALS_URI, status_code=status_code) + island_api_client.get_credentials_for_propagation() + + +def test_island_api_client_get_credentials_for_propagation__bad_json(): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(IslandAPIRequestFailedError): + m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad") + island_api_client.get_credentials_for_propagation()