UT: Fix faulty logic in test_multiple_status_queries

The test is checking that some calls return "in progress" and subsequent
calls return success. Using itertools.repeat allows all future calls to
return success.
This commit is contained in:
Mike Salvatore 2022-05-11 08:09:56 -04:00
parent 825c7b9ecf
commit d1d960abc2
1 changed files with 2 additions and 1 deletions

View File

@ -1,3 +1,4 @@
from itertools import chain, repeat
from unittest.mock import MagicMock
import pytest
@ -189,7 +190,7 @@ def test_multiple_status_queries(send_command_response, in_progress_response, su
aws_client = MagicMock()
aws_client.send_command = MagicMock(return_value=send_command_response)
aws_client.get_command_invocation = MagicMock(
side_effect=[in_progress_response, in_progress_response, success_response]
side_effect=chain([in_progress_response, in_progress_response], repeat(success_response))
)
command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP)