From 65c7ecbc30bcc066610bd5287f767d2c967496d9 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 21 Feb 2022 21:01:16 +0530
Subject: [PATCH] Test prompt drawing after screen is shrunk

---
 kitty_tests/__init__.py          |  2 ++
 kitty_tests/shell_integration.py | 18 +++++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/kitty_tests/__init__.py b/kitty_tests/__init__.py
index f3839e8ed3e..e31a5ca5e22 100644
--- a/kitty_tests/__init__.py
+++ b/kitty_tests/__init__.py
@@ -216,6 +216,8 @@ def wait_till(self, q, timeout=10):
             raise TimeoutError('The condition was not met')
 
     def set_window_size(self, rows=25, columns=80):
+        if hasattr(self, 'screen'):
+            self.screen.resize(rows, columns)
         x_pixels = columns * self.cell_width
         y_pixels = rows * self.cell_height
         s = struct.pack('HHHH', rows, columns, x_pixels, y_pixels)
diff --git a/kitty_tests/shell_integration.py b/kitty_tests/shell_integration.py
index be9cd9ab089..309fa66b602 100644
--- a/kitty_tests/shell_integration.py
+++ b/kitty_tests/shell_integration.py
@@ -70,7 +70,23 @@ def test_zsh_integration(self):
             self.ae(pty.callbacks.titlebuf, ['~'])
             pty.callbacks.clear()
             pty.send_cmd_to_child('mkdir test && ls -a')
-            pty.wait_till(lambda: pty.screen_contents().count(ps1) == 2)
+            pty.wait_till(lambda: pty.screen_contents().count(rps1) == 2)
             self.ae(pty.callbacks.titlebuf, ['mkdir test && ls -a', '~'])
             q = '\n'.join(str(pty.screen.line(i)) for i in range(1, pty.screen.cursor.y))
             self.ae(pty.last_cmd_output(), q)
+            # shrink the screen
+            pty.write_to_child(r'echo $COLUMNS')
+            pty.set_window_size(rows=20, columns=40)
+            q = ps1 + 'echo $COLUMNS' + ' ' * (40 - len(ps1) - len(rps1) - len('echo $COLUMNS')) + rps1
+            pty.process_input_from_child()
+
+            def redrawn():
+                q = pty.screen_contents()
+                return '$COLUMNS' in q and q.count(rps1) == 2 and q.count(ps1) == 2
+
+            pty.wait_till(redrawn)
+            self.ae(q, str(pty.screen.line(pty.screen.cursor.y)))
+            pty.write_to_child('\r')
+            pty.wait_till(lambda: pty.screen_contents().count(rps1) == 3)
+            self.ae('40', str(pty.screen.line(pty.screen.cursor.y - 1)))
+            self.ae(q, str(pty.screen.line(pty.screen.cursor.y - 2)))