def test_set_session(self, mock_session: Session, mock_random):
# Arrange
mock_random.return_value = 1
mock_session.return_value = mock_session
mock_session.cookies = MagicMock()
self.slave.ip_addr = "123.4567.8910"
# Act
self.slave.set_session(mock_session)
cookie = mock_session.cookies.set_cookie.call_args[0][0].value
# Assert
assert self.slave.ip_addr in cookie
def test_attempt_master_connection(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.status_code = 200
mock_session.get.return_value = mock_resp
mock_session.cookies = MagicMock()
# Act
self.slave.attempt_master_connection(5678)
# Assert
assert self.slave.session == mock_session
def test_connect_generic_exception(self, mock_session: Session, capsys):
# Arrange
mock_session.return_value = mock_session
mock_session.get.side_effect = Exception("Mock Error. Ignore Me!")
with pytest.raises(SystemExit):
# Act & Assert
assert self.slave.connect("hostname", "port")
def test_connect_connection_error(self, mock_session: Session):
# Arrange
mock_session.return_value = mock_session
mock_session.get.side_effect = ConnectionError
# Act
ret = self.slave.connect("hostname", "port")
# Assert
assert ret is None
def test_get_file_no_resp(self, mock_session: Session):
# Arrange
mock_session.return_value = mock_session
mock_session.get.return_value = False
self.slave.session = mock_session
# Act
success = self.slave.get_file("file_name")
# Assert
assert not success
def test_attempt_master_connection_none(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.status_code = 500
# Act
ret = self.slave.attempt_master_connection(5678)
# Assert
assert ret is None
def test_req_job_generic_exception(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.json.side_effect = Exception("Mock Error. Ignore Me!")
mock_session.get.return_value = mock_resp
self.slave.session = mock_session
# Act & Assert
with pytest.raises(Exception):
assert self.slave.req_job()
def test_connect_resp_not_ok(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.status_code = 500
mock_session.get.return_value = mock_resp
# Act
session = self.slave.connect("hostname", "port")
# Assert
assert session is None
def test_req_job_value_error(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.json.side_effect = ValueError
mock_session.get.return_value = mock_resp
self.slave.session = mock_session
# Act & Assert
with pytest.raises(Exception):
assert self.slave.req_job()
def test_get_file_decompress_exception(self, mock_session: Session,
mock_decompress):
# Arrange
mock_session.return_value = mock_session
mock_session.get.return_value = MagicMock()
self.slave.session = mock_session
mock_decompress.side_effect = CompressionException
# Act
success = self.slave.get_file("file_name")
# Assert
assert not success
def test_connect_invalid_json(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.status_code = 200
mock_resp.json.side_effect = ValueError
mock_session.get.return_value = mock_resp
with pytest.raises(Exception):
# Act & Assert
assert self.slave.connect("hostname", "port")
def test_connect_resp_200(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_resp.status_code = 200
mock_resp.json.return_value = "mock_json"
mock_session.get.return_value = mock_resp
# Act
session = self.slave.connect("hostname", "port")
# Assert
assert session == mock_session
def test_execute_tasks_failed(self, mock_session: Session,
mock_run_shell_command):
# Arrange
task_1: Task = Task(1, "", [], None, "result.txt", 'payload.txt')
task_1.message_type = TaskMessageType.TASK_RAW
tasks: List[Task] = [task_1]
mock_session.return_value = mock_session
self.slave.run_shell_command = MagicMock(return_value=1)
# Act
failed_tasks = self.slave.execute_tasks(tasks)
# Assert
assert failed_tasks == [task_1]
def test_req_task_endpoint(self, mock_session: Session):
# Arrange
mock_session.return_value = mock_session
mock_session.get.return_value = False
self.slave.host = "hostname"
self.slave.port = "port"
self.slave.job_id = 1234
self.slave.session = mock_session
expected_endpoint = f'http://{self.slave.host}:{self.slave.port}/{endpoints.GET_TASKS}/1234/1'
# Act
self.slave.req_tasks(1)
# Assert
mock_session.get.assert_called_with(expected_endpoint, timeout=5)
def test_req_job_endpoint(self, mock_session: Session,
mock_resp: Response):
# Arrange
mock_session.return_value = mock_session
mock_session.get.return_value = False
self.slave.host = "hostname"
self.slave.port = "port"
self.slave.session = mock_session
expected_endpoint = f'http://{self.slave.host}:{self.slave.port}/{endpoints.JOB}'
# Act
self.slave.req_job()
# Assert
mock_session.get.assert_called_with(expected_endpoint, timeout=5)
def test_req_task(self, mock_session: Session, mock_resp: Response):
# Arrange
expected_tasks: List[Task] = [Task(1, "", [], None, "", "")]
pickled_tasks = pickle_dumps(expected_tasks)
compressed_data = compress(pickled_tasks)
mock_resp.status_code = 69
mock_resp.content = compressed_data
mock_session.return_value = mock_session
mock_session.get.return_value = mock_resp
self.slave.session = mock_session
# Act
actual_tasks = self.slave.req_tasks(1)
# Assert
assert expected_tasks[0].task_id == actual_tasks[0].task_id
def test_req_job(self, mock_session: Session, mock_resp: Response):
# Arrange
mock_resp.json.return_value = '{"job_id": 1, "file_names": ["file_name"]}'
mock_session.return_value = mock_session
mock_session.get.return_value = mock_resp
self.slave.create_job_dir = MagicMock()
self.slave.get_file = MagicMock()
self.slave.handle_tasks = MagicMock()
self.slave.heartbeat = MagicMock()
self.slave.process_job = MagicMock(return_value=False)
self.slave.session = mock_session
# Act
self.slave.req_job()
# Assert
self.slave.get_file.assert_called_with('file_name')
def test_handle_tasks_save_files(self, mock_session: Session, mock_file):
# Arrange
expected_payload = "Test".encode()
task_1: Task = Task(1, "", [], expected_payload, "result.txt",
'payload.txt')
task_1.message_type = TaskMessageType.TASK_RAW
tasks: List[Task] = [task_1]
mock_session.return_value = mock_session
self.slave.req_tasks = MagicMock(return_value=tasks)
self.slave.save_processed_data = MagicMock()
self.slave.execute_tasks = MagicMock()
self.slave.session = mock_session
# Act
self.slave.handle_tasks(tasks)
# Assert
self.slave.execute_tasks.assert_called_with(tasks)
self.slave.save_processed_data.assert_called_with(
'payload.txt', expected_payload)
def test_get_file_endpoint(self, mock_session: Session,
mock_resp: Response, mock_decompress):
# Arrange
mock_session.return_value = mock_session
mock_resp.status_code = 200
mock_resp.json.return_value = "mock_json"
mock_session.get.return_value = mock_resp
job_id = "123"
file_name = "filename"
hostname = "hostname"
port = "port"
expected_endpoint = f'http://{hostname}:{port}/{endpoints.FILE}/{job_id}/{file_name}'
self.slave.host = hostname
self.slave.port = port
self.slave.session = mock_session
self.slave.job_id = job_id
self.slave.save_processed_data = MagicMock()
# Act
success = self.slave.get_file(file_name)
# Assert
mock_session.get.assert_called_with(expected_endpoint)
assert success