From 8d5a316d5462a98a957d3280e452ab3bf15ae11e Mon Sep 17 00:00:00 2001
From: Keannu Bernasol <noehbernasol0@gmail.com>
Date: Tue, 12 Sep 2023 21:36:28 +0800
Subject: [PATCH] Improved websocket consumer

---
 stude/config/asgi.py              | 13 ++++++++-----
 stude/student_status/consumers.py | 32 +++++++++++++++++++++++++------
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/stude/config/asgi.py b/stude/config/asgi.py
index e2bc07a..6384e94 100644
--- a/stude/config/asgi.py
+++ b/stude/config/asgi.py
@@ -1,6 +1,7 @@
 import os
 
 from channels.routing import ProtocolTypeRouter, URLRouter
+from channels.security.websocket import AllowedHostsOriginValidator
 from channels.auth import AuthMiddlewareStack
 from django.core.asgi import get_asgi_application
 import api.routing
@@ -13,9 +14,11 @@ django_asgi_app = get_asgi_application()
 
 application = ProtocolTypeRouter({
     "http": django_asgi_app,
-    'websocket': AuthMiddlewareStack(
-        URLRouter(
-            [re_path(r'ws/', URLRouter(api.routing.websocket_urlpatterns))]
-        )
-    ),
+    'websocket': AllowedHostsOriginValidator(
+        AuthMiddlewareStack(
+            URLRouter(
+                [re_path(r'ws/', URLRouter(api.routing.websocket_urlpatterns))]
+            )
+        ),
+    )
 })
diff --git a/stude/student_status/consumers.py b/stude/student_status/consumers.py
index 5d8a479..bbaf422 100644
--- a/stude/student_status/consumers.py
+++ b/stude/student_status/consumers.py
@@ -12,6 +12,12 @@ from djangochannelsrestframework.mixins import (
     RetrieveModelMixin,
 )
 from djangochannelsrestframework.permissions import IsAuthenticated
+from channels.layers import get_channel_layer
+from asgiref.sync import async_to_sync
+from django.contrib.gis.measure import Distance
+from django.contrib.gis.geos import fromstr
+from .models import StudentStatus
+from accounts.models import CustomUser
 
 
 class StudentStatusConsumer(
@@ -23,28 +29,42 @@ class StudentStatusConsumer(
     queryset = StudentStatus.objects.filter(active=True)
     serializer_class = StudentStatusSerializer
 
+    async def send_status_update(self, event):
+        data = event['data']
+        await self.send(text_data=json.dumps(data))
+
     async def websocket_connect(self, message):
-        # This method is called when the websocket is handshaking as part of the connection process.
+        await self.channel_layer.group_add('student_status_group', self.channel_name)
         await self.accept()
         self.send_updates_task = asyncio.create_task(self.send_updates())
 
     async def websocket_disconnect(self, message):
-        # This method is called when the WebSocket closes for any reason.
-        # Here we want to cancel our periodic task that sends updates
+        # ...
+        await self.channel_layer.group_discard('student_status_group', self.channel_name)
         self.send_updates_task.cancel()
 
     @database_sync_to_async
     def get_student_statuses(self):
         queryset = self.get_queryset()
+        user = self.scope['user']
         return StudentStatusSerializer(queryset, many=True).data
 
     async def send_updates(self):
+        channel_layer = get_channel_layer()
+
         while True:
             try:
+                print('attempting to get')
                 data = await self.get_student_statuses()
-                # print(f"Sending update: {data}") Debug
-                await self.send(text_data=json.dumps(data))
-                await asyncio.sleep(5)
+                print(f"Sending update: {data}")
+                await channel_layer.group_send(
+                    'student_status_group',
+                    {
+                        'type': 'send_status_update',
+                        'data': data,
+                    }
+                )
+                await asyncio.sleep(3)
             except Exception as e:
                 print(f"Exception in send_updates: {e}")
                 break  # Break the loop on error