From 8773240c0cb2cfd632b0efef0d808d6d63a4c938 Mon Sep 17 00:00:00 2001 From: jesopo Date: Wed, 11 Mar 2020 16:30:11 +0000 Subject: [PATCH] return Line objects successfully sent with StatefulEncoder.pop() --- irctokens/stateful.py | 5 ++++- test/stateful_encode.py | 18 +++++++++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/irctokens/stateful.py b/irctokens/stateful.py index e950e15..35edbdc 100644 --- a/irctokens/stateful.py +++ b/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)] diff --git a/test/stateful_encode.py b/test/stateful_encode.py index ce11fb7..2e8363b 100644 --- a/test/stateful_encode.py +++ b/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)