Add disconnected socket tests for conn.recv()

1 files changed, 46 insertions, 1 deletions

diff --git a/test/test_conn.py b/test/test_conn.py
index 71604cb..6a3b154 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -165,7 +165,52 @@ def test_can_send_more(conn):
     assert conn.can_send_more() is False
 
 
-def test_recv(socket, conn):
+def test_recv_disconnected():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(('127.0.0.1', 0))
+    port = sock.getsockname()[1]
+    sock.listen(5)
+
+    conn = BrokerConnection('127.0.0.1', port, socket.AF_INET)
+    timeout = time.time() + 1
+    while time.time() < timeout:
+        conn.connect()
+        if conn.connected():
+            break
+    else:
+        assert False, 'Connection attempt to local socket timed-out ?'
+
+    conn.send(MetadataRequest[0]([]))
+
+    # Disconnect server socket
+    sock.close()
+
+    # Attempt to receive should mark connection as disconnected
+    assert conn.connected()
+    conn.recv()
+    assert conn.disconnected()
+
+
+def test_recv_disconnected_too(_socket, conn):
+    conn.connect()
+    assert conn.connected()
+
+    req = MetadataRequest[0]([])
+    header = RequestHeader(req, client_id=conn.config['client_id'])
+    payload_bytes = len(header.encode()) + len(req.encode())
+    _socket.send.side_effect = [4, payload_bytes]
+    conn.send(req)
+
+    # Empty data on recv means the socket is disconnected
+    _socket.recv.return_value = b''
+
+    # Attempt to receive should mark connection as disconnected
+    assert conn.connected()
+    conn.recv()
+    assert conn.disconnected()
+
+
+def test_recv(_socket, conn):
     pass # TODO