From 4695f967288a52a480e42c5267ed512f18ed5856 Mon Sep 17 00:00:00 2001
From: Linn Mattsson <linma@cendio.se>
Date: Tue, 11 Oct 2022 07:31:48 +0000
Subject: Add new websocket class HttpWebSocket

This class acts as a glue between websocket and http functionality by
taking a 'request_handler' and using its functions for send_response(),
send_header() and end_headers().
---
 tests/test_websocketserver.py  | 69 ++++++++++++++++++++++++++++++++++++++++++
 websockify/websocketserver.py  | 21 +++++++++++--
 websockify/websockifyserver.py |  4 +--
 3 files changed, 90 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_websocketserver.py

diff --git a/tests/test_websocketserver.py b/tests/test_websocketserver.py
new file mode 100644
index 0000000..0e37e3d
--- /dev/null
+++ b/tests/test_websocketserver.py
@@ -0,0 +1,69 @@
+
+""" Unit tests for websocketserver """
+import unittest
+from unittest.mock import patch, MagicMock
+
+from websockify.websocketserver import HttpWebSocket
+
+
+class HttpWebSocketTest(unittest.TestCase):
+    @patch("websockify.websocketserver.WebSocket.__init__", autospec=True)
+    def test_constructor(self, websock):
+        # Given
+        req_obj = MagicMock()
+
+        # When
+        sock = HttpWebSocket(req_obj)
+
+        # Then
+        websock.assert_called_once_with(sock)
+        self.assertEqual(sock.request_handler, req_obj)
+
+    @patch("websockify.websocketserver.WebSocket.__init__", MagicMock(autospec=True))
+    def test_send_response(self):
+        # Given
+        req_obj = MagicMock()
+        sock = HttpWebSocket(req_obj)
+
+        # When
+        sock.send_response(200, "message")
+
+        # Then
+        req_obj.send_response.assert_called_once_with(200, "message")
+
+    @patch("websockify.websocketserver.WebSocket.__init__", MagicMock(autospec=True))
+    def test_send_response_default_message(self):
+        # Given
+        req_obj = MagicMock()
+        sock = HttpWebSocket(req_obj)
+
+        # When
+        sock.send_response(200)
+
+        # Then
+        req_obj.send_response.assert_called_once_with(200, None)
+
+    @patch("websockify.websocketserver.WebSocket.__init__", MagicMock(autospec=True))
+    def test_send_header(self):
+        # Given
+        req_obj = MagicMock()
+        sock = HttpWebSocket(req_obj)
+
+        # When
+        sock.send_header("keyword", "value")
+
+        # Then
+        req_obj.send_header.assert_called_once_with("keyword", "value")
+
+    @patch("websockify.websocketserver.WebSocket.__init__", MagicMock(autospec=True))
+    def test_end_headers(self):
+        # Given
+        req_obj = MagicMock()
+        sock = HttpWebSocket(req_obj)
+
+        # When
+        sock.end_headers()
+
+        # Then
+        req_obj.end_headers.assert_called_once_with()
+
diff --git a/websockify/websocketserver.py b/websockify/websocketserver.py
index 9088fe9..287f520 100644
--- a/websockify/websocketserver.py
+++ b/websockify/websocketserver.py
@@ -12,6 +12,23 @@ from http.server import BaseHTTPRequestHandler, HTTPServer
 
 from websockify.websocket import WebSocket, WebSocketWantReadError, WebSocketWantWriteError
 
+class HttpWebSocket(WebSocket):
+    """Class to glue websocket and http request functionality together"""
+    def __init__(self, request_handler):
+        super().__init__()
+
+        self.request_handler = request_handler
+
+    def send_response(self, code, message=None):
+        self.request_handler.send_response(code, message)
+
+    def send_header(self, keyword, value):
+        self.request_handler.send_header(keyword, value)
+
+    def end_headers(self):
+        self.request_handler.end_headers()
+
+
 class WebSocketRequestHandlerMixIn:
     """WebSocket request handler mix-in class
 
@@ -25,7 +42,7 @@ class WebSocketRequestHandlerMixIn:
     use for the WebSocket connection.
     """
 
-    SocketClass = WebSocket
+    SocketClass = HttpWebSocket
 
     def handle_one_request(self):
         """Extended request handler
@@ -59,7 +76,7 @@ class WebSocketRequestHandlerMixIn:
         The WebSocket object will then replace the request object and
         handle_websocket() will be called.
         """
-        websocket = self.SocketClass()
+        websocket = self.SocketClass(self)
         try:
             websocket.accept(self.request, self.headers)
         except Exception:
diff --git a/websockify/websockifyserver.py b/websockify/websockifyserver.py
index 0199e42..3e13fa1 100644
--- a/websockify/websockifyserver.py
+++ b/websockify/websockifyserver.py
@@ -29,10 +29,10 @@ if sys.platform == 'win32':
     # make sockets pickle-able/inheritable
     import multiprocessing.reduction
 
-from websockify.websocket import WebSocket, WebSocketWantReadError, WebSocketWantWriteError
+from websockify.websocket import WebSocketWantReadError, WebSocketWantWriteError
 from websockify.websocketserver import WebSocketRequestHandlerMixIn
 
-class CompatibleWebSocket(WebSocket):
+class CompatibleWebSocket(WebSocketRequestHandlerMixIn.SocketClass):
     def select_subprotocol(self, protocols):
         # Handle old websockify clients that still specify a sub-protocol
         if 'binary' in protocols:
-- 
cgit v1.2.1