gh-110673: test_pty raises on short write

vstinner · vstinner · commit b5ebb45fba82 · 2023-10-11T22:22:20.000+02:00
Add write_all() helper function to test_pty to raise an exception on
short write: if os.writes() does not write all bytes. It should not
happen for a PTY.
diff --git a/Lib/test/test_pty.py b/Lib/test/test_pty.py
@@ -76,6 +76,15 @@ def expectedFailureIfStdinIsTTY(fun):
         pass
     return fun
 
+
+def write_all(fd, data):
+    written = os.write(fd, data)
+    if written != len(data):
+        # gh-73256, gh-110673: It should never happen, but check just in case
+        raise Exception(f"short write: os.write({fd}, {len(data)} bytes) "
+                        f"wrote {written} bytes")
+
+
 # Marginal testing of pty suite. Cannot do extensive 'do or fail' testing
 # because pty code is not too portable.
 class PtyTest(unittest.TestCase):
@@ -170,14 +179,14 @@ def test_openpty(self):
             os.set_blocking(master_fd, blocking)
 
         debug("Writing to slave_fd")
-        os.write(slave_fd, TEST_STRING_1)
+        write_all(slave_fd, TEST_STRING_1)
         s1 = _readline(master_fd)
         self.assertEqual(b'I wish to buy a fish license.\n',
                          normalize_output(s1))
 
         debug("Writing chunked output")
-        os.write(slave_fd, TEST_STRING_2[:5])
-        os.write(slave_fd, TEST_STRING_2[5:])
+        write_all(slave_fd, TEST_STRING_2[:5])
+        write_all(slave_fd, TEST_STRING_2[5:])
         s2 = _readline(master_fd)
         self.assertEqual(b'For my pet fish, Eric.\n', normalize_output(s2))
 
@@ -360,8 +369,8 @@ def test__copy_to_each(self):
         masters = [s.fileno() for s in socketpair]
 
         # Feed data.  Smaller than PIPEBUF.  These writes will not block.
-        os.write(masters[1], b'from master')
-        os.write(write_to_stdin_fd, b'from stdin')
+        write_all(masters[1], b'from master')
+        write_all(write_to_stdin_fd, b'from stdin')
 
         # Expect three select calls, the last one will cause IndexError
         pty.select = self._mock_select
