Basic websocket support

Author: kvalev

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~=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),
     ))