return Line objects successfully sent with StatefulEncoder.pop()

irctokens/stateful.py

@@ -25,12 +25,15 @@ class StatefulDecoder(object):
 class StatefulEncoder(object):
     def __init__(self):
         self._buffer = b""
+        self._buffered_lines: typing.List[Line] = []
 
     def pending(self) -> bytes:
         return self._buffer
 
     def push(self, line: Line) -> bytes:
         self._buffer += f"{line.format()}\r\n".encode("utf8")
-        return self._buffer
+        self._buffered_lines.append(line)
     def pop(self, byte_count: int):
+        sent = self._buffer[:byte_count].count("\n")
         self._buffer = self._buffer[byte_count:]
+        return [self._buffered_lines.pop(0) for _ in range(sent)]

test/stateful_encode.py

@@ -9,9 +9,25 @@ class TestPush(unittest.TestCase):
         self.assertEqual(e.pending(), b"PRIVMSG #channel hello\r\n")
 
 class TestPop(unittest.TestCase):
-    def test(self):
+    def test_partial(self):
         e = irctokens.StatefulEncoder()
         line = irctokens.tokenise("PRIVMSG #channel hello")
         e.push(line)
         e.pop(len(b"PRIVMSG #channel hello"))
         self.assertEqual(e.pending(), b"\r\n")
+
+    def test_returned(self):
+        e = irctokens.StatefulEncoder()
+        line = irctokens.tokenise("PRIVMSG #channel hello")
+        e.push(line)
+        e.push(line)
+        lines = e.pop(len(b"PRIVMSG #channel hello\r\nPRIVMSG"))
+        self.assertEqual(len(lines), 1)
+        self.assertEqual(lines[0], line)
+
+    def test_none_returned(self):
+        e = irctokens.StatefulEncoder()
+        line = irctokens.tokenise("PRIVMSG #channel hello")
+        e.push(line)
+        lines = e.pop(1)
+        self.assertEqual(len(lines), 0)