connectionManager.py

from fastapi import WebSocket

from datetime import datetime
from typing import List

class Connection: 
    websocket: WebSocket
    connection_time: datetime

    def __init__(self, websocket: WebSocket, connection_time: datetime):
        self.websocket = websocket
        self.connection_time = connection_time

class ConnectionManager:
    timeout = 60 * 5 # 5 minutes

    def __init__(self):
        self.active_connections: List[Connection] = []

    async def connect(self, websocket: WebSocket):
        print('Connecting')
        await websocket.accept()
        # Add connection time and websocket to active connections
        self.active_connections.append(Connection(websocket=websocket, connection_time=datetime.now()))

    def isConnected(self, websocket: WebSocket):
        for connection in self.active_connections:
            if connection.websocket == websocket:
                return True
        return False
  
    def shouldDisconnect(self, websocket: WebSocket):
        for connection in self.active_connections:
            if connection.websocket == websocket:
                if (datetime.now() - connection.connection_time).total_seconds() > self.timeout:
                    print('Disconnecting...')
                    return True
        return False

    async def receive_json(self, websocket: WebSocket):
        if not self.isConnected(websocket):
          return None
        print('Receiving...')
        data = await websocket.receive_json()
        print('Received')
        return data

    def disconnect(self, websocket: WebSocket):
        print('Disconnecting...')
        for connection in self.active_connections:
            if connection.websocket == websocket:
                self.active_connections.remove(connection)
                return True
        return False

    async def send_json(self, json, websocket: WebSocket):
        print('Sending JSON...')
        # Only send the message if the connection is still active
        if self.isConnected(websocket):
          await websocket.send_json(json)