diff --git a/promptshell/main.py b/promptshell/main.py index a8dfb90..484f3d0 100644 --- a/promptshell/main.py +++ b/promptshell/main.py @@ -32,7 +32,7 @@ def main(): if len(prompt) + len(user_input) > columns: print() # Move to the next line if input is too long - if user_input.lower() == 'quit': + if user_input.lower() in ['quit', 'exit']: # Allow both 'quit' and 'exit' commands to terminate the REPL session print(format_text('red', bold=True) + "\nTerminating..." + reset_format()) break diff --git a/tests/test_repl.py b/tests/test_repl.py new file mode 100644 index 0000000..7a62f1d --- /dev/null +++ b/tests/test_repl.py @@ -0,0 +1,23 @@ +import pytest +from io import StringIO +from unittest.mock import patch +from promptshell.main import main # Adjust the import based on your project structure +import os + +@patch('os.get_terminal_size', return_value=os.terminal_size((80, 24))) +@patch('builtins.input', side_effect=['exit']) +@patch('sys.stdout', new_callable=StringIO) +def test_exit_command(mock_stdout, mock_input, mock_terminal_size): + """Test that the 'exit' command terminates the REPL session.""" + main() + output = mock_stdout.getvalue() + assert "Terminating..." in output # Check for termination message + +@patch('os.get_terminal_size', return_value=os.terminal_size((80, 24))) +@patch('builtins.input', side_effect=['quit']) +@patch('sys.stdout', new_callable=StringIO) +def test_quit_command(mock_stdout, mock_input, mock_terminal_size): + """Test that the 'quit' command terminates the REPL session.""" + main() + output = mock_stdout.getvalue() + assert "Terminating..." in output # Check for termination message