Basic websocket support

pook/interceptors/__init__.py

@@ -1,4 +1,5 @@
 import sys
+from .websocket import WebSocketInterceptor
 from .urllib3 import Urllib3Interceptor
 from .http import HTTPClientInterceptor
 from .base import BaseInterceptor
@@ -7,6 +8,7 @@
 __all__ = (
  'interceptors', 'add', 'get',
  'BaseInterceptor',
+ 'WebSocketInterceptor',
  'Urllib3Interceptor',
  'HTTPClientInterceptor',
  'AIOHTTPInterceptor',
@@ -15,7 +17,8 @@
 # Store built-in interceptors in pook.
 interceptors = [
  Urllib3Interceptor,
- HTTPClientInterceptor
+ HTTPClientInterceptor,
+ WebSocketInterceptor
 ]
 
 # Import aiohttp in modern Python runtimes

pook/interceptors/websocket.py

@@ -0,0 +1,115 @@
+from ..request import Request
+from .base import BaseInterceptor
+
+# Support Python 2/3
+try:
+ import mock
+except Exception:
+ from unittest import mock
+
+
+class MockFrame(object):
+ def __init__(self, opcode, data, fin=1):
+ self.opcode = opcode
+ self.data = data
+ self.fin = fin
+
+
+class MockSocket(object):
+ def fileno(self):
+ return 0
+
+ def gettimeout(self):
+ return 0
+
+
+class MockHandshakeResponse(object):
+
+ def __init__(self, status, headers, subprotocol):
+ self.status = status
+ self.headers = headers
+ self.subprotocol = subprotocol
+
+
+class WebSocketInterceptor(BaseInterceptor):
+
+ def _handshake(self, sock, hostname, port, resource, **options):
+ return MockHandshakeResponse(200, {}, "")
+
+ def _connect(self, url, options, proxy, socket):
+ req = Request()
+ req.headers = {} # TODO
+ req.url = url
+
+ # TODO does this work multithreaded?!?
+ self._mock = self.engine.match(req)
+ if not self._mock:
+ # we cannot forward, as we have mocked away the connection
+ raise ValueError(
+ "Request to '%s' could not be matched or forwarded" % url
+ )
+
+ # make the body always a list to simplify our lives
+ body = self._mock._response._body
+ if not isinstance(body, list):
+ self._mock._response._body = [body]
+
+ sock = MockSocket()
+ addr = ("hostname", "port", "resource")
+
+ return sock, addr
+
+ def _send(self, data):
+ # mock as if all data has been sent
+ return len(data)
+
+ def _recv_frame(self):
+ # alias
+ body = self._mock._response._body
+
+ idx = getattr(self, "_data_index", 0)
+ if len(body) >= idx:
+ # close frame
+ return MockFrame(0x8, None)
+
+ # data frame
+ self._data_index = idx + 1
+ return MockFrame(0x1, body[idx].encode("utf-8"), 1)
+
+ def _patch(self, path, handler):
+ try:
+ # Create a new patcher for Urllib3 urlopen function
+ # used as entry point for all the HTTP communications
+ patcher = mock.patch(path, handler)
+ # Retrieve original patched function that we might need for real
+ # networking
+ # request = patcher.get_original()[0]
+ # Start patching function calls
+ patcher.start()
+ except Exception:
+ # Exceptions may accur due to missing package
+ # Ignore all the exceptions for now
+ pass
+ else:
+ self.patchers.append(patcher)
+
+ def activate(self):
+ """
+ Activates the traffic interceptor.
+ This method must be implemented by any interceptor.
+ """
+ patches = [
+ ('websocket._core.connect', self._connect),
+ ('websocket._core.handshake', self._handshake),
+ ('websocket.WebSocket._send', self._send),
+ ('websocket.WebSocket.recv_frame', self._recv_frame),
+ ]
+
+ [self._patch(path, handler) for path, handler in patches]
+
+ def disable(self):
+ """
+ Disables the traffic interceptor.
+ This method must be implemented by any interceptor.
+ """
+ [patch.stop() for patch in self.patchers]

pook/matchers/url.py

@@ -10,8 +10,8 @@
 else: # Python 3
  from urllib.parse import urlparse
 
-# URI protocol test regular expression
-protoregex = re.compile('^http[s]?://', re.IGNORECASE)
+# URI protocol test regular expression (without group capture)
+protoregex = re.compile('^(?:ws|http)[s]?://', re.IGNORECASE)
 
 
 class URLMatcher(BaseMatcher):

requirements-dev.txt

@@ -12,3 +12,4 @@ urllib3>=1.24.2
 bumpversion~=0.5.3
 aiohttp~=1.1.5 ; python_version >= '3.4.2'
 mocket~=1.6.0
+websocket-client~=0.57.0

tests/unit/interceptors/websocket_test.py

@@ -0,0 +1,13 @@
+from websocket import WebSocket
+import pook
+
+
+@pook.on
+def test_websocket():
+ (pook.get('ws://some-non-existing.org')
+ .reply(204)
+ .body('test'))
+
+ socket = WebSocket()
+ socket.connect('ws://some-non-existing.org/', header=['x-custom: header'])
+ socket.send('test')

tests/unit/matchers/url_test.py

@@ -32,6 +32,8 @@ def test_url_matcher_urlparse():
  ('http://foo.com/foo/bar', 'http://foo.com/foo/bar', True),
  ('http://foo.com/foo/bar/baz', 'http://foo.com/foo/bar/baz', True),
  ('http://foo.com/foo?x=y&z=w', 'http://foo.com/foo?x=y&z=w', True),
+ ('ws://foo.com', 'ws://foo.com', True),
+ ('wss://foo.com', 'wss://foo.com', True),
 
  # Invalid cases
  ('http://foo.com', 'http://bar.com', False),
@@ -41,6 +43,7 @@ def test_url_matcher_urlparse():
  ('http://foo.com/foo/bar', 'http://foo.com/bar/foo', False),
  ('http://foo.com/foo/bar/baz', 'http://foo.com/baz/bar/foo', False),
  ('http://foo.com/foo?x=y&z=w', 'http://foo.com/foo?x=x&y=y', False),
+ ('ws://foo.com', 'wss://foo.com', False),
  ))