bootstrap-python-fastapi/tests/http_app/routes/ws/test_chat.py at 844fceee0b0704e1e09a82f15094475a6fcf21d5 · febus982/bootstrap-python-fastapi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from unittest.mock import patch

import pytest
from fastapi.testclient import TestClient

from http_app.routes.ws.chat import ConnectionManager


@pytest.fixture
def test_client(testapp):
    return TestClient(testapp)


@pytest.fixture
def connection_manager():
    return ConnectionManager()


def test_websocket_connection(test_client):
    manager = ConnectionManager()
    with patch("http_app.routes.ws.chat.manager", manager):
        with test_client.websocket_connect("/ws/chat/1") as websocket:
            websocket.send_text("Hello!")
            data = websocket.receive_text()

            assert data == "You wrote: Hello!"
            broadcast = websocket.receive_text()
            assert broadcast == "Client #1 says: Hello!"


def test_multiple_clients(test_client):
    manager = ConnectionManager()
    with patch("http_app.routes.ws.chat.manager", manager):
        with test_client.websocket_connect("/ws/chat/1") as websocket1:
            with test_client.websocket_connect("/ws/chat/2") as websocket2:
                # Client 1 sends message
                websocket1.send_text("Hello from client 1")

                # Client 1 receives personal message
                data1 = websocket1.receive_text()
                assert data1 == "You wrote: Hello from client 1"

                # Both clients receive broadcast
                broadcast1 = websocket1.receive_text()
                broadcast2 = websocket2.receive_text()
                assert broadcast1 == "Client #1 says: Hello from client 1"
                assert broadcast2 == "Client #1 says: Hello from client 1"


def test_client_disconnect(test_client):
    manager = ConnectionManager()
    with patch("http_app.routes.ws.chat.manager", manager):
        with test_client.websocket_connect("/ws/chat/1") as websocket1:
            with test_client.websocket_connect("/ws/chat/2") as websocket2:
                # Close first client
                websocket1.close()

                # Second client should receive disconnect message
                disconnect_message = websocket2.receive_text()
                assert disconnect_message == "Client #1 left the chat"


async def test_connection_manager():
    manager = ConnectionManager()

    # Mock WebSocket for testing
    class MockWebSocket:
        def __init__(self):
            self.received_messages = []

        async def accept(self):
            pass

        async def send_text(self, message: str):
            self.received_messages.append(message)

    # Test connect
    websocket = MockWebSocket()
    await manager.connect(websocket)
    assert len(manager.active_connections) == 1

    # Test personal message
    await manager.send_personal_message("test message", websocket)
    assert websocket.received_messages[-1] == "test message"

    # Test broadcast
    await manager.broadcast("broadcast message")
    assert websocket.received_messages[-1] == "broadcast message"

    # Test disconnect
    manager.disconnect(websocket)
    assert len(manager.active_connections) == 0