mirror of https://github.com/OWASP/Nettacker.git
26 lines
761 B
Python
26 lines
761 B
Python
from io import StringIO
|
|
from unittest.mock import patch
|
|
|
|
from nettacker.core.die import die_success, die_failure
|
|
from nettacker.logger import TerminalCodes
|
|
|
|
|
|
@patch("sys.stdout", new_callable=StringIO)
|
|
@patch("sys.exit")
|
|
def test_die_success(mock_exit, mock_stdout):
|
|
reset_code = TerminalCodes.RESET.value
|
|
die_success()
|
|
success_message = mock_stdout.getvalue()
|
|
assert reset_code in success_message
|
|
mock_exit.assert_called_once_with(0)
|
|
|
|
|
|
@patch("sys.stdout", new_callable=StringIO)
|
|
@patch("sys.exit")
|
|
def test_die_failure(mock_exit, mock_stdout):
|
|
test_message = "Test error message"
|
|
die_failure(test_message)
|
|
error_message = mock_stdout.getvalue()
|
|
assert test_message in error_message
|
|
mock_exit.assert_called_once_with(1)
|