UT: Add tests for HTTPIslandAPIClient.get_agent_binary()

This commit is contained in:
Shreya Malviya 2022-09-20 17:28:22 +05:30
parent 30cf360e98
commit e8ecaa2169
1 changed files with 37 additions and 0 deletions

View File

@ -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)