Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -532,9 +532,9 @@ async def connect(self):
raise ValueError(
"Please install aiohttp library to use websocket transport."
)
except OSError:
except OSError as e:
await self.session.close()
raise ConnectionError('Client Session Closed')
raise ConnectionError('Websocket connection closed: %r' % e) from e
Comment thread
swathipil marked this conversation as resolved.

async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argument
"""Read exactly n bytes from the peer."""
Expand All @@ -558,6 +558,9 @@ async def _read(self, n, buffer=None, **kwargs): # pylint: disable=unused-argum
return view
except asyncio.TimeoutError as te:
raise ConnectionError('Receive timed out (%s)' % te)
except OSError as e:
await self.session.close()
raise ConnectionError('Websocket connection closed: %r' % e) from e

async def close(self):
"""Do any preliminary work in shutting down the connection."""
Expand All @@ -575,3 +578,6 @@ async def write(self, s):
await self.ws.send_bytes(s)
except asyncio.TimeoutError as te:
raise ConnectionError('Send timed out (%s)' % te)
except OSError as e:
await self.session.close()
raise ConnectionError('Websocket connection closed: %r' % e) from e
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
import pytest
import asyncio
from unittest.mock import patch

import aiohttp
from azure.eventhub._pyamqp.aio._transport_async import WebSocketTransportAsync


# class WebsocketException(unittest.TestCase):
async def test_websocket_aiohttp_exception():
Comment thread
swathipil marked this conversation as resolved.
with patch.object(aiohttp.ClientSession,'ws_connect', side_effect=aiohttp.ClientOSError):
transport = WebSocketTransportAsync(host="my_host")
with pytest.raises(ConnectionError):
await transport.connect()