Fixes #1517

tests/test_raw_websocket.py

@@ -143,6 +143,53 @@ class ParseCloseTest(unittest.TestCase):
         self.assertEqual(RawWebSocket._parse_close(b''), (None, ''))
 
 
+class CloseTest(unittest.IsolatedAsyncioTestCase):
+    async def test_peer_close_still_closes_transport_without_another_close_frame(self):
+        reader = asyncio.StreamReader()
+        reader.feed_data(_raw_frame(RawWebSocket.OP_CLOSE, b'\x03\xe8'))
+        writer = Mock(drain=AsyncMock(), wait_closed=AsyncMock())
+        ws = RawWebSocket(reader, writer)
+
+        self.assertIsNone(await ws.recv())
+        await ws.close()
+
+        writer.close.assert_called_once()
+        writer.wait_closed.assert_awaited_once()
+        self.assertEqual(writer.write.call_count, 1)
+
+    async def test_repeated_close_sends_only_one_close_frame(self):
+        writer = Mock(drain=AsyncMock(), wait_closed=AsyncMock())
+        ws = RawWebSocket(asyncio.StreamReader(), writer)
+
+        await ws.close()
+        await ws.close()
+
+        self.assertTrue(ws._closed)
+        self.assertEqual(writer.write.call_count, 1)
+        writer.close.assert_called()
+
+    async def test_cancelling_close_during_drain_still_closes_transport(self):
+        draining = asyncio.Event()
+
+        async def drain():
+            draining.set()
+            await asyncio.Future()
+
+        writer = Mock(drain=AsyncMock(side_effect=drain), wait_closed=AsyncMock())
+        ws = RawWebSocket(asyncio.StreamReader(), writer)
+        task = asyncio.create_task(ws.close())
+        try:
+            await asyncio.wait_for(draining.wait(), 1)
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+            writer.close.assert_called_once()
+            writer.wait_closed.assert_awaited_once()
+        finally:
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
 class HandshakeErrorTest(unittest.TestCase):
     def test_redirect_status_codes(self):
         for code in (301, 302, 303, 307, 308):