Store http code and headers in WSServerHandshakeError

aio-libs · Dec 25, 2015 · 1419f0c · 1419f0c
1 parent 62bd199
commit 1419f0c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/aiohttp/client.py b/aiohttp/client.py
@@ -273,20 +273,32 @@ def _ws_connect(self, url, *,
         try:
             # check handshake
             if resp.status != 101:
-                raise WSServerHandshakeError('Invalid response status')
+                raise WSServerHandshakeError(
+                    message='Invalid response status',
+                    code=resp.status,
+                    headers=resp.headers)
 
             if resp.headers.get(hdrs.UPGRADE, '').lower() != 'websocket':
-                raise WSServerHandshakeError('Invalid upgrade header')
+                raise WSServerHandshakeError(
+                    message='Invalid upgrade header',
+                    code=resp.status,
+                    headers=resp.headers)
 
             if resp.headers.get(hdrs.CONNECTION, '').lower() != 'upgrade':
-                raise WSServerHandshakeError('Invalid connection header')
+                raise WSServerHandshakeError(
+                    message='Invalid connection header',
+                    code=resp.status,
+                    headers=resp.headers)
 
             # key calculation
             key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, '')
             match = base64.b64encode(
                 hashlib.sha1(sec_key + WS_KEY).digest()).decode()
             if key != match:
-                raise WSServerHandshakeError('Invalid challenge response')
+                raise WSServerHandshakeError(
+                    message='Invalid challenge response',
+                    code=resp.status,
+                    headers=resp.headers)
 
             # websocket protocol
             protocol = None

diff --git a/aiohttp/errors.py b/aiohttp/errors.py
@@ -96,9 +96,6 @@ def __init__(self, *, code=None, message='', headers=None):
 class WSServerHandshakeError(HttpProcessingError):
     """websocket server handshake error."""
 
-    def __init__(self, message, *, headers=None):
-        super().__init__(message=message, headers=headers)
-
 
 class HttpProxyError(HttpProcessingError):
     """Http proxy error.