Fix a couple of missed failing tests

sanic-org · Sep 23, 2021 · cb495ac · cb495ac
1 parent 37d462a
commit cb495ac
Show file tree

Hide file tree

Showing 6 changed files with 148 additions and 24 deletions.
diff --git a/examples/websocket_bench.html b/examples/websocket_bench.html
@@ -0,0 +1,47 @@
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>WebSocket benchmark</title>
+    </head>
+    <body>
+        <script>
+            function make_msg(length) {
+                var result = '';
+                var characters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+                var clen = characters.length;
+                for ( var i = 0; i < length; i++ ) {
+                  result += characters.charAt(Math.floor(Math.random() * clen));
+               }
+               return result;
+            }
+            var ws = new WebSocket('ws://' + document.domain + ':' + location.port + '/bench');
+            var msg = make_msg(16384);
+            var started = false;
+            function send_msg() {
+                if (ws.readyState === ws.OPEN) {
+                    ws.send(msg);
+                    window.setTimeout(send_msg, 1);
+                }
+            }
+            function run_bench() {
+                console.log("first msg seen, starting bench...");
+                window.setTimeout(send_msg, 0);
+            }
+            function stop_bench() {
+                ws.close(1000);
+            }
+            ws.onmessage = function(event) {
+                console.log("Received "+ event.data);
+                if (started === false) {
+                    started = true;
+                    window.setTimeout(run_bench, 0);
+                    window.setTimeout(stop_bench, 40000);
+                }
+            };
+            ws.onclose = function(event) {
+                console.log("Websocket connection closed.");
+            }
+
+        </script>
+    </body>
+</html>
diff --git a/examples/websocket_bench.py b/examples/websocket_bench.py
@@ -0,0 +1,70 @@
+import asyncio
+import logging
+from sanic import Sanic
+from sanic.response import file
+from sanic.log import error_logger
+import time
+error_logger.setLevel(logging.INFO)
+app = Sanic(__name__)
+
+@app.route('/')
+async def index(request):
+    return await file('websocket_bench.html')
+
+
+@app.websocket('/bench')
+async def bench_p_time(request, ws):
+    i = 0
+    bytes_total = 0
+    start = 0
+    end_time = 0
+    started = False
+    await ws.send("1")
+    while started is False or (end_time > time.time()):
+        i += 1
+        in_data = await ws.recv()
+        if started is False:
+            del in_data
+            error_logger.info("received first data: starting benchmark now..")
+            started = True
+            start = time.time()
+            end_time = start + 30.0
+            continue
+        bytes_total += len(in_data)
+        del in_data
+    end = time.time()
+    elapsed = end - start
+    error_logger.info("Done. Took {} seconds".format(elapsed))
+    error_logger.info("{} bytes in 30 seconds = {}".format(bytes_total, (bytes_total/30.0)))
+
+@app.websocket('/benchp')
+async def bench_p_time(request, ws):
+    i = 0
+    bytes_total = 0
+    real_start = 0
+    start_ptime = 0
+    end_ptime = 0
+    started = False
+    await ws.send("1")
+    while started is False or (end_ptime > time.process_time()):
+        i += 1
+        in_data = await ws.recv()
+        if started is False:
+            del in_data
+            error_logger.info("received first data: starting benchmark now..")
+            started = True
+            real_start = time.time()
+            start_ptime = time.process_time()
+            end_ptime = start_ptime + 30.0
+            continue
+        bytes_total += len(in_data)
+        del in_data
+    real_end = time.time()
+    elapsed = real_end - real_start
+    error_logger.info("Done. Took {} seconds".format(elapsed))
+    error_logger.info("{} bytes in 30 seconds = {}".format(bytes_total, (bytes_total/30.0)))
+
+
+if __name__ == '__main__':
+    app.run(host="0.0.0.0", port=8000, debug=False, auto_reload=False)
+
diff --git a/sanic/server/websockets/impl.py b/sanic/server/websockets/impl.py
@@ -334,8 +334,8 @@ def end_connection(self, code=1000, reason=""):
         if self.connection.state == OPEN:
             data_to_send = self.connection.data_to_send()
             self.connection.send_close(code, reason)
+            data_to_send.extend(self.connection.data_to_send())
             try:
-                data_to_send.extend(self.connection.data_to_send())
                 while (
                     len(data_to_send)
                     and self.io_proto
@@ -346,15 +346,18 @@ def end_connection(self, code=1000, reason=""):
             except Exception:
                 # sending close frames may fail if the
                 # transport closes during this period
+                # But that doesn't matter at this point
                 ...
         if (
             self.data_finished_fut is not None
             and not self.data_finished_fut.done()
         ):
-            # We have a graceful auto-closer. Use it to close the connection.
+            # We have the ability to signal the auto-closer
+            # try to trigger it to auto-close the connection
             self.data_finished_fut.cancel()
             self.data_finished_fut = None
         if self.auto_closer_task is None or self.auto_closer_task.done():
+            # Auto-closer is not running, do force disconnect
             return self._force_disconnect()
 
     async def auto_close_connection(self) -> None:

diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
@@ -16,6 +16,7 @@
     abort,
 )
 from sanic.response import text
+from websockets.version import version as websockets_version
 
 
 class SanicExceptionTestException(Exception):
@@ -260,9 +261,14 @@ async def feed(request, ws):
 
     with caplog.at_level(logging.INFO):
         app.test_client.websocket("/feed")
-
-    assert caplog.record_tuples[1][0] == "sanic.error"
-    assert caplog.record_tuples[1][1] == logging.ERROR
+    # Websockets v10.0 and above output an additional
+    # INFO message when a ws connection is accepted
+    ws_version_parts = websockets_version.split(".")
+    ws_major = int(ws_version_parts[0])
+    record_index = 2 if ws_major >= 10 else 1
+    assert caplog.record_tuples[record_index][0] == "sanic.error"
+    assert caplog.record_tuples[record_index][1] == logging.ERROR
     assert (
-        "Exception occurred while handling uri:" in caplog.record_tuples[1][2]
+        "Exception occurred while handling uri:"
+        in caplog.record_tuples[record_index][2]
     )
diff --git a/tests/test_routes.py b/tests/test_routes.py
@@ -671,31 +671,29 @@ async def check(request):
     assert response.json["set"]
 
 
-def test_websocket_route_with_subprotocols(app):
+@pytest.mark.parametrize(
+    "subprotocols,expected",
+    (
+        (["one"], "one"),
+        (["three", "one"], "one"),
+        (["tree"], None),
+        (None, None),
+    ),
+)
+def test_websocket_route_with_subprotocols(app, subprotocols, expected):
     results = []
 
     @app.websocket("/ws", subprotocols=["zero", "one", "two", "three"])
     async def handler(request, ws):
-        results.append(ws.subprotocol)
+        nonlocal results
+        results = ws.subprotocol
         assert ws.subprotocol is not None
 
-    _, response = SanicTestClient(app).websocket("/ws", subprotocols=["one"])
-    assert response.opened is True
-    assert results == ["one"]
-
     _, response = SanicTestClient(app).websocket(
-        "/ws", subprotocols=["three", "one"]
+        "/ws", subprotocols=subprotocols
     )
     assert response.opened is True
-    assert results == ["one", "one"]
-
-    _, response = SanicTestClient(app).websocket("/ws", subprotocols=["tree"])
-    assert response.opened is True
-    assert results == ["one", "one", None]
-
-    _, response = SanicTestClient(app).websocket("/ws")
-    assert response.opened is True
-    assert results == ["one", "one", None, None]
+    assert results == expected
 
 
 @pytest.mark.parametrize("strict_slashes", [True, False, None])

diff --git a/tests/test_worker.py b/tests/test_worker.py
@@ -175,7 +175,7 @@ def test_worker_close(worker):
     worker.wsgi = mock.Mock()
     conn = mock.Mock()
     conn.websocket = mock.Mock()
-    conn.websocket.close = mock.Mock(wraps=_a_noop)
+    conn.websocket.fail_connection = mock.Mock(wraps=_a_noop)
     worker.connections = set([conn])
     worker.log = mock.Mock()
     worker.loop = loop
@@ -190,5 +190,5 @@ def test_worker_close(worker):
     loop.run_until_complete(_close)
 
     assert worker.signal.stopped
-    assert conn.websocket.close.called
+    assert conn.websocket.fail_connection.called
     assert len(worker.servers) == 0