Add tests for long stdio messages, fix await logic

Author: krassowski

python_packages/jupyter_lsp/jupyter_lsp/stdio.py

@@ -93,12 +93,14 @@ async def read(self) -> None:
                     self.wake()
 
                 IOLoop.current().add_callback(self.queue.put_nowait, message)
-            except Exception:  # pragma: no cover
-                self.log.exception("%s couldn't enqueue message: %s", self, message)
+            except Exception as e:  # pragma: no cover
+                self.log.exception(
+                    "%s couldn't enqueue message: %s (%s)", self, message, e
+                )
                 await self.sleep()
 
-    def _read_content(
-        self, length: int, max_parts=100, max_empties=50
+    async def _read_content(
+        self, length: int, max_parts=500, max_empties=50
     ) -> Optional[bytes]:
         """Read the full length of the message unless exceeding max_parts or
            max_empties empty reads occur.
@@ -119,13 +121,16 @@ def _read_content(
         raw = None
         raw_parts: List[bytes] = []
         received_size = 0
-        while (
-            received_size < length and len(raw_parts) < max_parts and max_empties > 0
-        ):
-            part = self.stream.read(length)
+        while received_size < length and len(raw_parts) < max_parts and max_empties > 0:
+            part = None
+            try:
+                part = self.stream.read(length)
+            except OSError:  # pragma: no cover
+                pass
             if part is None:
                 max_empties -= 1
-                continue  # pragma: no cover
+                await self.sleep()
+                continue
             received_size += len(part)
             raw_parts.append(part)
 
@@ -135,6 +140,7 @@ def _read_content(
                 self.log.warning(
                     f"Readout and content-length mismatch:" f" {len(raw)} vs {length}"
                 )
+
         return raw
 
     async def read_one(self) -> Text:
@@ -152,24 +158,15 @@ async def read_one(self) -> Text:
             content_length = int(headers.get("content-length", "0"))
 
             if content_length:
-                raw = None
-                retries = 5
-                while raw is None and retries:
-                    try:
-                        raw = self._read_content(length=content_length)
-                    except OSError:  # pragma: no cover
-                        raw = None
-                    if raw is None:  # pragma: no cover
-                        self.log.warning(
-                            "%s failed to read message of length %s",
-                            self,
-                            content_length,
-                        )
-                        await self.sleep()
-                        retries -= 1
-                    else:
-                        message = raw.decode("utf-8").strip()
-                        break
+                raw = await self._read_content(length=content_length)
+                if raw is not None:
+                    message = raw.decode("utf-8").strip()
+                else:  # pragma: no cover
+                    self.log.warning(
+                        "%s failed to read message of length %s",
+                        self,
+                        content_length,
+                    )
 
         return message
 

python_packages/jupyter_lsp/jupyter_lsp/tests/test_stdio.py

@@ -0,0 +1,81 @@
+import asyncio
+import subprocess
+import time
+from threading import Thread
+
+import pytest
+from tornado.queues import Queue
+
+from jupyter_lsp.stdio import LspStdIoReader
+
+WRITER_TEMPLATE = """
+from time import sleep
+
+print('Content-Length: {length}')
+print()
+
+for repeat in range({repeats}):
+    sleep({interval})
+    print('{message}', end='')
+print()
+"""
+
+
+class CommunicatorSpawner:
+    def __init__(self, tmpdir):
+        self.tmpdir = tmpdir
+
+    def spawn_writer(self, message: str, repeats: int = 1, interval=None):
+        length = len(message) * repeats
+        commands_file = self.tmpdir / "writer.py"
+        commands_file.write(
+            WRITER_TEMPLATE.format(
+                length=length, repeats=repeats, interval=interval or 0, message=message
+            )
+        )
+        return subprocess.Popen(
+            ["python", "-u", commands_file.realpath()],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+        )
+
+
+@pytest.fixture
+def communicator_spawner(tmpdir):
+    return CommunicatorSpawner(tmpdir)
+
+
+def communicate_and_close(process, wait=1):
+    def communicate_and_close():
+        time.sleep(wait)
+        process.communicate()
+
+    thread = Thread(target=communicate_and_close)
+    thread.start()
+
+
+@pytest.mark.parametrize(
+    "message,repeats,interval",
+    [
+        ["short", 1, None],
+        ["ab" * 10_0000, 1, None],
+        ["ab", 2, 0.01],
+        ["ab", 45, 0.01],
+    ],
+    ids=["short", "long", "intermittent", "intensive-intermittent"],
+)
+@pytest.mark.asyncio
+async def test_reader(message, repeats, interval, communicator_spawner):
+    queue = Queue()
+
+    process = communicator_spawner.spawn_writer(
+        message=message, repeats=repeats, interval=interval
+    )
+    reader = LspStdIoReader(stream=process.stdout, queue=queue)
+    timeout = 2 + (interval or 1) * repeats * 2
+
+    communicate_and_close(process)
+    await asyncio.wait_for(reader.read(), timeout=timeout)
+
+    result = queue.get_nowait()
+    assert result == message * repeats