Work on testing infrastructure for DnD kitten

kitty_tests/__init__.py

@@ -336,7 +336,7 @@ class PTY:
 
     def __init__(
         self, argv=None, rows=25, columns=80, scrollback=100, cell_width=10, cell_height=20,
-        cwd=None, env=None, stdin_fd=None, stdout_fd=None, needs_da1=True,
+        cwd=None, env=None, stdin_fd=None, stdout_fd=None, needs_da1=True, window_id=0,
     ):
         self.is_child = False
         if isinstance(argv, str):
@@ -377,7 +377,7 @@ class PTY:
         self.set_window_size(rows=rows, columns=columns)
         self.needs_da1 = needs_da1
         self.callbacks = Callbacks(self)
-        self.screen = Screen(self.callbacks, rows, columns, scrollback, cell_width, cell_height, 0, self.callbacks)
+        self.screen = Screen(self.callbacks, rows, columns, scrollback, cell_width, cell_height, window_id, self.callbacks)
         self.received_bytes = b''
 
     def reset_termios_state(self):
@@ -393,6 +393,12 @@ class PTY:
         s = termios.tcgetattr(self.master_fd)
         return True if s[3] & termios.ECHO else False
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.__del__()
+
     def __del__(self):
         if not self.is_child:
             if hasattr(self, 'master_fd'):

kitty_tests/dnd_kitten.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+# License: GPLv3 Copyright: 2026, Kovid Goyal <kovid at kovidgoyal.net>
+
+import tempfile
+from base64 import standard_b64encode
+
+from kitty.constants import kitten_exe
+from kitty.fast_data_types import DND_CODE, dnd_set_test_write_func, dnd_test_cleanup_fake_window, dnd_test_create_fake_window, dnd_test_probe_state
+
+from . import PTY, BaseTest
+from .dnd import WriteCapture
+
+
+class Capture(WriteCapture):
+
+    def __call__(self, window_id: int, data: bytes) -> None:
+        self.pty.write(data)
+
+
+class TestDnDKitten(BaseTest):
+
+    def setUp(self):
+        capture = Capture()
+        dnd_set_test_write_func(capture)
+        os_window_id, window_id = dnd_test_create_fake_window()
+        capture.window_id = window_id
+        capture.os_window_id = os_window_id
+        self.capture = capture
+        self.test_dir = self.enterContext(tempfile.TemporaryDirectory())
+        self.messages_from_kitten = ''
+
+    def send_dnd_command_to_kitten(self, payload=b'', as_base64=False, flush=False, **metadata):
+        header = f'\x1b]{DND_CODE};'
+        for k, v in metadata.items():
+            header = header + f'{k}={v}:'
+        self.pty.write_to_child(header.encode())
+        if not payload:
+            self.pty.write_to_child(b'\x1b\\', flush=flush)
+            return
+        if isinstance(payload, str):
+            payload = payload.encode()
+        payload = memoryview(standard_b64encode(payload) if as_base64 else payload)
+        for i in range(0, len(payload), 4096):
+            end = i + 4096
+            is_last = end >= len(payload)
+            chunk = payload[i:min(i+4096, len(payload))]
+            if i == 0:
+                self.pty.write_to_child(f'm={0 if is_last else 1};'.encode())
+            else:
+                self.pty.write_to_child(f'\x1b]{DND_CODE};m={0 if is_last else 1};'.encode())
+            self.pty.write_to_child(chunk)
+            self.pty.write_to_child(b'\x1b\\', flush=is_last and flush)
+
+    def finish_setup(self, cmd=None):
+        cmd = cmd or [kitten_exe(), 'dnd']
+        self.pty = self.enterContext(PTY(argv=cmd, rows=25, columns=80, window_id=self.capture.window_id))
+        self.pty.callbacks.printbuf = self
+        self.screen = self.pty.screen
+        self.pty.wait_till(lambda: bool(self.pty.callbacks.titlebuf))
+        self.assertFalse(self.probe_state('drop_is_remote_client'))
+
+    def append(self, text):
+        self.messages_from_kitten += text
+
+    def wait_for_responses(self, *responses, timeout=10):
+        q = '\n'.join(responses)
+        def wait_till():
+            return q == self.messages_from_kitten.strip()
+        self.pty.wait_till(wait_till, timeout, lambda: f'Responses so far: {self.messages_from_kitten!r}')
+        self.messages_from_kitten = ''
+
+    def probe_state(self, which: str):
+        return dnd_test_probe_state(self.capture.window_id, which)
+
+    def tearDown(self):
+        dnd_set_test_write_func(None)
+        dnd_test_cleanup_fake_window(self.capture.os_window_id)
+        del self.capture
+        del self.screen
+        del self.pty
+
+    def test_dnd_kitten_drop(self):
+        self.finish_setup()