From d6252bcb16b38485fb6f7127fa74900b45fc3da1 Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Mon, 9 Apr 2018 20:02:55 +0000
Subject: [PATCH] Add more socket-related testing helpers.

---
 tests/helpers/socket.py | 63 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/tests/helpers/socket.py b/tests/helpers/socket.py
index 94463d23..4d2fd0f1 100644
--- a/tests/helpers/socket.py
+++ b/tests/helpers/socket.py
@@ -1,10 +1,34 @@
 from __future__ import absolute_import
 
 from collections import namedtuple
+import contextlib
+import errno
 
 import ptvsd.socket as _ptvsd
 
 
+try:
+    BrokenPipeError  # noqa
+except NameError:
+    class BrokenPipeError(Exception):
+        pass
+
+
+NOT_CONNECTED = (
+    errno.ENOTCONN,
+    errno.EBADF,
+)
+
+CLOSED = (
+    errno.EPIPE,
+    errno.ESHUTDOWN,
+)
+
+EOF = NOT_CONNECTED + CLOSED
+
+
+# TODO: Add timeouts to the functions.
+
 def create_server(address):
     """Return a server socket after binding."""
     host, port = address
@@ -53,6 +77,45 @@ def bind(address):
     return connect, remote
 
 
+def recv_as_read(sock):
+    """Return a wrapper ardoung sock.read that arises EOFError when closed."""
+    def read(numbytes, _recv=sock.recv):
+        try:
+            return _recv(numbytes)
+        except BrokenPipeError:
+            raise EOFError
+        except OSError as exc:
+            if exc.errno in EOF:
+                raise EOFError
+            raise
+    return read
+
+
+def send_as_write(sock):
+    """Return a wrapper ardoung sock.send that arises EOFError when closed."""
+    def write(data, _send=sock.send):
+        try:
+            return _send(data)
+        except BrokenPipeError:
+            raise EOFError
+        except OSError as exc:
+            if exc.errno in EOF:
+                raise EOFError
+            raise
+    return write
+
+
+@contextlib.contextmanager
+def timeout(sock, timeout):
+    """A context manager that sets a timeout on blocking socket ops."""
+    orig = sock.gettimeout()
+    sock.settimeout(timeout)
+    try:
+        yield
+    finally:
+        sock.settimeout(orig)
+
+
 def close(sock):
     """Shutdown and close the socket."""
     _ptvsd.close_socket(sock)