Added fake websocket client

Author: Joel Collins

tests/conftest.py

@@ -7,9 +7,39 @@
 from labthings.server.labthing import LabThing
 from labthings.server.view import View
 
+from werkzeug.test import EnvironBuilder
 from flask.testing import FlaskClient
 
 
+class FakeWebsocket:
+    def __init__(self, message: str, recieve_once=True):
+        self.message = message
+        self.responses = []
+        self.closed = False
+        self.recieve_once = recieve_once
+
+        # I mean screw whoever is responsible for this having to be a thing...
+        self.receive = self.recieve
+
+    def recieve(self):
+        # Get message
+        message_to_send = self.message
+        # If only sending a message to the server once
+        if self.recieve_once:
+            # Clear our message
+            self.message = None
+        return message_to_send
+
+    @property
+    def response(self):
+        return self.responses[-1]
+
+    def send(self, response):
+        self.responses.append(response)
+        self.closed = True
+        return response
+
+
 class JsonClient(FlaskClient):
     def open(self, *args, **kwargs):
         kwargs.setdefault(
@@ -20,6 +50,57 @@ def open(self, *args, **kwargs):
         return super().open(*args, **kwargs)
 
 
+class SocketClient(FlaskClient):
+    def __init__(self, app, response_wrapper, *args, **kwargs):
+        super().__init__(app, response_wrapper, *args, **kwargs)
+        self.app = app
+        self.response_wrapper = response_wrapper
+        self.socket = FakeWebsocket(message=None)
+        self.environ_base = {
+            "HTTP_UPGRADE": "websocket",
+            "wsgi.websocket": self.socket,
+        }
+
+    def connect(self, *args, message=None, **kwargs):
+        kwargs.setdefault("environ_overrides", {})[
+            "flask._preserve_context"
+        ] = self.preserve_context
+        kwargs.setdefault("environ_base", self.environ_base)
+        builder = EnvironBuilder(*args, **kwargs)
+
+        try:
+            environ = builder.get_environ()
+        finally:
+            builder.close()
+
+        self.socket.message = message
+
+        with self.app.app_context():
+            run_wsgi_app(self.app, environ)
+
+        # Once the connection has been closed, return responses
+        return self.socket.responses
+
+
+def run_wsgi_app(app, environ, buffered=False):
+    response = []
+    buffer = []
+
+    def start_response(status, headers, exc_info=None):
+        if exc_info:
+            try:
+                raise exc_info[1].with_traceback(exc_info[2])
+            finally:
+                exc_info = None
+        response[:] = [status, headers]
+        return buffer.append
+
+    # Return value from the wsgi_app call
+    # In the case of our SocketMiddleware, will return []
+    app_rv = app(environ, start_response)
+    return app_rv
+
+
 @pytest.fixture
 def empty_view_cls():
     class EmptyViewClass(View):
@@ -146,6 +227,12 @@ def client(app):
     return app.test_client()
 
 
+@pytest.fixture
+def ws_client(app):
+    app.test_client_class = SocketClient
+    return app.test_client()
+
+
 @pytest.fixture
 def thing_client(thing):
     thing.app.test_client_class = JsonClient
@@ -167,28 +254,6 @@ def extensions_path(app):
     return os.path.join(os.path.dirname(__file__), "extensions")
 
 
-class FakeWebsocket:
-    def __init__(self, message: str, recieve_once=True):
-        self.message = message
-        self.response = None
-        self.closed = False
-        self.recieve_once = recieve_once
-
-    def receive(self):
-        # Get message
-        message_to_send = self.message
-        # If only sending a message to the server once
-        if self.recieve_once:
-            # Clear our message
-            self.message = None
-        return message_to_send
-
-    def send(self, response):
-        self.response = response
-        self.closed = True
-        return response
-
-
 @pytest.fixture
 def fake_websocket():
     """