Modified async_alert_str to also get passed the number of terminal columns and added unit tests for it

3 files changed, 19 insertions, 10 deletions

diff --git a/cmd2/ansi.py b/cmd2/ansi.py
index c0d8017a..f1b2def8 100644
--- a/cmd2/ansi.py
+++ b/cmd2/ansi.py
@@ -188,21 +188,17 @@ style_warning = functools.partial(style, fg='bright_yellow')
 style_error = functools.partial(style, fg='bright_red')
 
 
-def async_alert_str(*, prompt: str, line: str, cursor_offset: int, alert_msg: str) -> str:
+def async_alert_str(*, terminal_columns: int, prompt: str, line: str, cursor_offset: int, alert_msg: str) -> str:
     """Calculate the desired string, including ANSI escape codes, for displaying an asynchronous alert message.
 
+    :param terminal_columns: terminal width (number of columns)
     :param prompt: prompt that is displayed on the current line
     :param line: current contents of the Readline line buffer
     :param cursor_offset: the offset of the current cursor position within line
     :param alert_msg: the message to display to the user
     :return: the correct string so that the alert message appears to the user to be printed above the current line.
     """
-    import shutil
     from colorama import Cursor
-
-    # Get the size of the terminal
-    terminal_size = shutil.get_terminal_size()
-
     # Split the prompt lines since it can contain newline characters.
     prompt_lines = prompt.splitlines()
 
@@ -211,7 +207,7 @@ def async_alert_str(*, prompt: str, line: str, cursor_offset: int, alert_msg: st
     num_prompt_terminal_lines = 0
     for line in prompt_lines[:-1]:
         line_width = ansi_safe_wcswidth(line)
-        num_prompt_terminal_lines += int(line_width / terminal_size.columns) + 1
+        num_prompt_terminal_lines += int(line_width / terminal_columns) + 1
 
     # Now calculate how many terminal lines are take up by the input
     last_prompt_line = prompt_lines[-1]
@@ -219,13 +215,13 @@ def async_alert_str(*, prompt: str, line: str, cursor_offset: int, alert_msg: st
 
     input_width = last_prompt_line_width + ansi_safe_wcswidth(line)
 
-    num_input_terminal_lines = int(input_width / terminal_size.columns) + 1
+    num_input_terminal_lines = int(input_width / terminal_columns) + 1
 
     # Get the cursor's offset from the beginning of the first input line
     cursor_input_offset = last_prompt_line_width + cursor_offset
 
     # Calculate what input line the cursor is on
-    cursor_input_line = int(cursor_input_offset / terminal_size.columns) + 1
+    cursor_input_line = int(cursor_input_offset / terminal_columns) + 1
 
     # Create a string that when printed will clear all input lines and display the alert
     terminal_str = ''
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 7efa6849..a10219b1 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -3830,7 +3830,9 @@ class Cmd(cmd.Cmd):
                     update_terminal = True
 
             if update_terminal:
-                terminal_str = ansi.async_alert_str(prompt=current_prompt, line=readline.get_line_buffer(),
+                import shutil
+                terminal_str = ansi.async_alert_str(terminal_columns=shutil.get_terminal_size().columns,
+                                                    prompt=current_prompt, line=readline.get_line_buffer(),
                                                     cursor_offset=rl_get_point(), alert_msg=alert_msg)
                 if rl_type == RlType.GNU:
                     sys.stderr.write(terminal_str)
diff --git a/tests/test_ansi.py b/tests/test_ansi.py
index 53c917d5..75e5ba35 100644
--- a/tests/test_ansi.py
+++ b/tests/test_ansi.py
@@ -99,3 +99,14 @@ def test_set_title_str():
     BEL = '\007'
     title = HELLO_WORLD
     assert ansi.set_title_str(title) == OSC + '2;' + title + BEL
+
+
+@pytest.mark.parametrize('cols, prompt, line, cursor, msg, expected', [
+    (127, '(Cmd) ', 'help his', 12, ansi.style('Hello World!', fg='magenta'), '\x1b[2K\r\x1b[35mHello World!\x1b[39m'),
+    (127, '\n(Cmd) ', 'help ', 5, 'foo', '\x1b[2K\x1b[1A\x1b[2K\rfoo'),
+    (10, '(Cmd) ', 'help history of the american republic', 4, 'boo', '\x1b[3B\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K\x1b[1A\x1b[2K\rboo')
+])
+def test_async_alert_str(cols, prompt, line, cursor, msg, expected):
+    alert_str = ansi.async_alert_str(terminal_columns=cols, prompt=prompt, line=line, cursor_offset=cursor,
+                                     alert_msg=msg)
+    assert  alert_str == expected