p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

UT: Add test for should_agent_stop

This commit is contained in:
Kekoa Kaaikala 2022-09-20 16:30:24 +00:00
parent 44d8dbeb5c
commit fd08212763
2 changed files with 53 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
monkey/infection_monkey/island_api_client/http_island_api_client.py
View File

@ -35,6 +35,14 @@ def handle_island_errors(fn):
def decorated(*args, **kwargs): def decorated(*args, **kwargs):
try: try:
return fn(*args, **kwargs) return fn(*args, **kwargs)
except (
IslandAPIConnectionError,
IslandAPIError,
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
) as e:
raise e
except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err: except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err:
raise IslandAPIConnectionError(err) raise IslandAPIConnectionError(err)
except requests.exceptions.HTTPError as err: except requests.exceptions.HTTPError as err:

45
monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py
View File

@ -39,6 +39,7 @@ 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}" ISLAND_GET_AGENT_BINARY_URI = f"https://{SERVER}/api/agent-binaries/{WINDOWS}"
ISLAND_SEND_EVENTS_URI = f"https://{SERVER}/api/agent-events" ISLAND_SEND_EVENTS_URI = f"https://{SERVER}/api/agent-events"
ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents" ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents"
ISLAND_AGENT_STOP_URI = f"https://{SERVER}/api/monkey-control/needs-to-stop/{AGENT_ID}"
class Event1(AbstractAgentEvent): class Event1(AbstractAgentEvent):
@ -318,3 +319,47 @@ def test_island_api_client_register_agent__status_code(status_code, expected_err
with pytest.raises(expected_error): with pytest.raises(expected_error):
m.post(ISLAND_REGISTER_AGENT_URI, status_code=status_code) m.post(ISLAND_REGISTER_AGENT_URI, status_code=status_code)
island_api_client.register_agent(AGENT_REGISTRATION) island_api_client.register_agent(AGENT_REGISTRATION)
@pytest.mark.parametrize(
"actual_error, expected_error",
[
(requests.exceptions.ConnectionError, IslandAPIConnectionError),
(TimeoutError, IslandAPITimeoutError),
],
)
def test_island_api_client__should_agent_stop(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_AGENT_STOP_URI, exc=actual_error)
island_api_client.should_agent_stop(AGENT_ID)
@pytest.mark.parametrize(
"status_code, expected_error",
[
(401, IslandAPIRequestError),
(501, IslandAPIRequestFailedError),
],
)
def test_island_api_client_should_agent_stop__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_AGENT_STOP_URI, status_code=status_code)
island_api_client.should_agent_stop(AGENT_ID)
def test_island_api_client_should_agent_stop__bad_json():
with requests_mock.Mocker() as m:
m.get(ISLAND_URI)
island_api_client = HTTPIslandAPIClient(SERVER)
with pytest.raises(IslandAPIRequestFailedError):
m.get(ISLAND_AGENT_STOP_URI, content=b"bad")
island_api_client.should_agent_stop(AGENT_ID)