Add websocket unittests and put in users app

.github/workflows/ci.yml

@@ -70,6 +70,8 @@ jobs:
       fail-fast: false
       matrix:
         script:
+          - name: With Async
+            args: "use_async=y"
           - name: With Celery
             args: "use_celery=y use_compressor=y"
           - name: With Gulp

hooks/post_gen_project.py

@@ -125,7 +125,13 @@ def remove_celery_files():
 def remove_async_files():
     file_names = [
         os.path.join("config", "asgi.py"),
-        os.path.join("config", "websocket.py"),
+        os.path.join("{{cookiecutter.project_slug}}", "users", "websocket.py"),
+        os.path.join(
+            "{{cookiecutter.project_slug}}", "users", "tests", "async_server.py"
+        ),
+        os.path.join(
+            "{{cookiecutter.project_slug}}", "users", "tests", "test_socket.py"
+        ),
     ]
     for file_name in file_names:
         os.remove(file_name)

{{cookiecutter.project_slug}}/config/asgi.py

@@ -28,7 +28,7 @@
 # application = HelloWorldApplication(application)
 
 # Import websocket application here, so apps from django_application are loaded first
-from config.websocket import websocket_application  # noqa isort:skip
+from {{ cookiecutter.project_slug }}.users.websocket import websocket_application  # noqa isort:skip
 
 
 async def application(scope, receive, send):

{{cookiecutter.project_slug}}/config/settings/test.py

@@ -15,6 +15,11 @@
 # https://docs.djangoproject.com/en/dev/ref/settings/#test-runner
 TEST_RUNNER = "django.test.runner.DiscoverRunner"
 
+{%- if cookiecutter.use_async == 'y' %}
+# Needed for socket testing that also needs HTTP to go along with it
+ALLOWED_HOSTS = ["localhost", "0.0.0.0", "127.0.0.1"]
+{%- endif %}
+
 # PASSWORDS
 # ------------------------------------------------------------------------------
 # https://docs.djangoproject.com/en/dev/ref/settings/#password-hashers

{{cookiecutter.project_slug}}/requirements/base.txt

@@ -25,6 +25,7 @@ flower==1.0.0  # https://github.com/mher/flower
 {%- endif %}
 {%- if cookiecutter.use_async == 'y' %}
 uvicorn[standard]==0.16.0  # https://github.com/encode/uvicorn
+websockets==10.1  # https://github.com/aaugustin/websockets
 {%- endif %}
 
 # Django

{{cookiecutter.project_slug}}/requirements/local.txt

@@ -15,11 +15,14 @@ watchgod==0.7  # https://github.com/samuelcolvin/watchgod
 # ------------------------------------------------------------------------------
 mypy==0.930  # https://github.com/python/mypy
 django-stubs==1.9.0  # https://github.com/typeddjango/django-stubs
-pytest==6.2.5  # https://github.com/pytest-dev/pytest
-pytest-sugar==0.9.4  # https://github.com/Frozenball/pytest-sugar
 {%- if cookiecutter.use_drf == "y" %}
 djangorestframework-stubs==1.4.0  # https://github.com/typeddjango/djangorestframework-stubs
 {%- endif %}
+pytest==6.2.5  # https://github.com/pytest-dev/pytest
+pytest-sugar==0.9.4  # https://github.com/Frozenball/pytest-sugar
+{%- if cookiecutter.use_async == 'y' %}
+pytest-timeout==1.4.2  # https://github.com/pytest-dev/pytest-timeout/
+{%- endif %}
 
 # Documentation
 # ------------------------------------------------------------------------------

{{cookiecutter.project_slug}}/{{cookiecutter.project_slug}}/users/tests/async_server.py

@@ -0,0 +1,40 @@
+import asyncio
+import functools
+import threading
+import time
+from contextlib import contextmanager
+
+from uvicorn.config import Config
+from uvicorn.main import ServerState
+from uvicorn.protocols.http.h11_impl import H11Protocol
+from uvicorn.protocols.websockets.websockets_impl import WebSocketProtocol
+
+
+def run_loop(loop):
+    loop.run_forever()
+    loop.close()
+
+
+@contextmanager
+def run_server(app, path="/"):
+    asyncio.set_event_loop(None)
+    loop = asyncio.new_event_loop()
+    config = Config(app=app, ws=WebSocketProtocol)
+    server_state = ServerState()
+    protocol = functools.partial(H11Protocol, config=config, server_state=server_state)
+    create_server_task = loop.create_server(protocol, host="127.0.0.1")
+    server = loop.run_until_complete(create_server_task)
+    port = server.sockets[0].getsockname()[1]  # type: ignore
+    url = "ws://127.0.0.1:{port}{path}".format(port=port, path=path)
+    try:
+        # Run the event loop in a new thread.
+        thread = threading.Thread(target=run_loop, args=[loop])
+        thread.start()
+        # Return the contextmanager state.
+        yield url
+    finally:
+        # Close the loop from our main thread.
+        while server_state.tasks:
+            time.sleep(0.01)
+        loop.call_soon_threadsafe(loop.stop)
+        thread.join()

{{cookiecutter.project_slug}}/{{cookiecutter.project_slug}}/users/tests/test_socket.py

@@ -0,0 +1,33 @@
+from asyncio import new_event_loop
+
+import pytest
+from websockets import connect
+
+from {{cookiecutter.project_slug}}.users.tests.async_server import run_server
+from {{cookiecutter.project_slug}}.users.websocket import websocket_application as app
+
+
+def test_accept_connection():
+    async def open_connection(url):
+        async with connect(url) as websocket:
+            return websocket.open
+
+    with run_server(app) as _url:
+        loop = new_event_loop()
+        is_open = loop.run_until_complete(open_connection(_url))
+        assert is_open
+        loop.close()
+
+
+@pytest.mark.timeout(10)
+def test_ping():
+    async def open_connection(url):
+        async with connect(url) as websocket:
+            await websocket.send("ping")
+            return await websocket.recv()
+
+    with run_server(app) as _url:
+        loop = new_event_loop()
+        received_message = loop.run_until_complete(open_connection(_url))
+        assert received_message == "pong"
+        loop.close()

{{cookiecutter.project_slug}}/config/websocket.py → {{cookiecutter.project_slug}}/{{cookiecutter.project_slug}}/users/websocket.py

@@ -1,13 +1,18 @@
 async def websocket_application(scope, receive, send):
+    event = await receive()
+    if event["type"] == "websocket.connect":
+        # TODO Add authentication by reading scope
+        #  and getting sessionid from cookie
+        await send({"type": "websocket.accept"})
+    else:
+        await send({"type": "websocket.close"})
+        return
+
     while True:
         event = await receive()
-
-        if event["type"] == "websocket.connect":
-            await send({"type": "websocket.accept"})
-
         if event["type"] == "websocket.disconnect":
             break
 
         if event["type"] == "websocket.receive":
             if event["text"] == "ping":
-                await send({"type": "websocket.send", "text": "pong!"})
+                await send({"type": "websocket.send", "text": "pong"})