diff --git a/spectator/writer/line_buffer.py b/spectator/writer/line_buffer.py index 35d1599..8f13e0a 100644 --- a/spectator/writer/line_buffer.py +++ b/spectator/writer/line_buffer.py @@ -2,7 +2,7 @@ class LineBuffer: - """Line buffer for Writers.""" + """Line buffer for Writers, sized in UTF-8 bytes including newline separators.""" def __init__(self, max_size: int) -> None: self._size: int = 0 @@ -14,8 +14,11 @@ def __len__(self) -> int: def append(self, line: str) -> bool: """Return value indicates whether the buffer should be flushed.""" + line_size = len(line) if line.isascii() else len(line.encode("utf-8")) + if self._lines: + self._size += 1 # Newline inserted between lines by flush(). self._lines.append(line) - self._size += len(line) + self._size += line_size return self._size >= self._max_size def flush(self) -> str: diff --git a/tests/writer/test_line_buffer.py b/tests/writer/test_line_buffer.py index 98373c1..6e7228f 100644 --- a/tests/writer/test_line_buffer.py +++ b/tests/writer/test_line_buffer.py @@ -14,7 +14,7 @@ def test_all_methods(self): self.assertFalse(should_flush) should_flush = buffer.append("e") - self.assertEqual(5, len(buffer)) + self.assertEqual(6, len(buffer)) self.assertTrue(should_flush) self.assertEqual("abcd\ne", buffer.flush()) @@ -26,3 +26,39 @@ def test_all_methods(self): self.assertEqual("abcdefghij", buffer.flush()) self.assertEqual(0, len(buffer)) + + def test_flush_threshold_includes_newline_separators(self): + buffer = LineBuffer(7) + self.assertFalse(buffer.append("abc")) + self.assertTrue(buffer.append("def")) + self.assertEqual(7, len(buffer)) + self.assertEqual("abc\ndef", buffer.flush()) + + def test_size_matches_utf8_payload(self): + buffer = LineBuffer(6) + self.assertFalse(buffer.append("é")) + self.assertEqual(2, len(buffer)) + self.assertTrue(buffer.append("€")) + self.assertEqual(6, len(buffer)) + payload = buffer.flush() + self.assertEqual(6, len(payload.encode("utf-8"))) + self.assertEqual(0, len(buffer)) + self.assertFalse(buffer.append("abc")) + self.assertEqual(3, len(buffer)) + + def test_empty_lines_include_separators(self): + buffer = LineBuffer(1) + self.assertFalse(buffer.append("")) + self.assertTrue(buffer.append("")) + self.assertEqual(1, len(buffer)) + self.assertEqual("\n", buffer.flush()) + + def test_invalid_utf8_leaves_buffer_unchanged(self): + buffer = LineBuffer(10) + self.assertFalse(buffer.append("foo")) + with self.assertRaises(UnicodeEncodeError): + buffer.append("\ud800") + self.assertEqual(3, len(buffer)) + self.assertFalse(buffer.append("bar")) + self.assertEqual(7, len(buffer)) + self.assertEqual("foo\nbar", buffer.flush()) diff --git a/tests/writer/test_socket_writer_udp.py b/tests/writer/test_socket_writer_udp.py index 36523bd..278955f 100644 --- a/tests/writer/test_socket_writer_udp.py +++ b/tests/writer/test_socket_writer_udp.py @@ -1,6 +1,7 @@ import threading import unittest from contextlib import closing +from unittest.mock import patch from spectator import Config, SocketWriter from ..udp_server import UdpServer @@ -34,10 +35,15 @@ def test_udp_with_buffer_flushes_on_close(self) -> None: def test_udp_close_closes_socket_on_encode_error(self) -> None: with closing(UdpServer()) as server: w = SocketWriter(Config(server.address(), buffer_size=1024)) - w.write("\ud800") + w.write("foo") sock = w._sock # pylint: disable=protected-access assert sock is not None - self.assertRaises(UnicodeEncodeError, w.close) + # Invalid text is now rejected by write(); simulate an encoding failure + # during close to retain coverage of the socket cleanup guarantee. + buffer = w._buffer # pylint: disable=protected-access + assert buffer is not None + with patch.object(buffer, "flush", return_value="\ud800"): + self.assertRaises(UnicodeEncodeError, w.close) self.assertEqual(-1, sock.fileno()) w.write("foo") self.assertIsNone(w._sock) # pylint: disable=protected-access