async def test_too_many_rooms_created(app: Starlette) -> None:
for i in range(MAX_ROOMS_PER_TEN_MINUTES):
async with emulated_client.connect(app, f'/{uuid4()}') as client:
await client.receive_json()
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(app, f'/{uuid4()}') as client:
await client.receive_json()
assert e.value.code == ERR_TOO_MANY_ROOMS_CREATED
async def test_bypass_room_create_rate_limit(app: Starlette) -> None:
headers = {BYPASS_RATE_LIMIT_HEADER: TEST_BYPASS_RATE_LIMIT_KEY}
for i in range(MAX_ROOMS_PER_TEN_MINUTES):
async with emulated_client.connect(
app, f'/{uuid4()}', headers=headers
) as client:
await client.receive_json()
async with emulated_client.connect(app, f'/{uuid4()}', headers=headers) as client:
assert await client.receive_json() == {'type': 'connected', 'data': []}
async def test_bypass_room_create_rate_limit_invalid_key(app: Starlette) -> None:
headers = {BYPASS_RATE_LIMIT_HEADER: "invalid-key"}
for i in range(MAX_ROOMS_PER_TEN_MINUTES):
async with emulated_client.connect(
app, f'/{uuid4()}', headers=headers
) as client:
await client.receive_json()
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(
app, f'/{uuid4()}', headers=headers
) as client:
await client.receive_json()
assert e.value.code == ERR_TOO_MANY_ROOMS_CREATED
async def connect(i: int) -> EmulatedClient:
ip = f'127.0.0.{i}' if unique_ips else '127.0.0.1'
room_id = uuid4() if unique_rooms else ROOM_ID
return await stack.enter_async_context(
emulated_client.connect(
app, f'/{room_id}', client_ip=ip, headers=headers
)
)
async def test_bypass_max_connections_rate_limit(app: Starlette) -> None:
headers = {BYPASS_RATE_LIMIT_HEADER: TEST_BYPASS_RATE_LIMIT_KEY}
async with num_active_connections(
app, MAX_CONNECTIONS_PER_USER, unique_ips=False, headers=headers
):
async with emulated_client.connect(
app, f'/{uuid4()}', headers=headers
) as client:
assert await client.receive_json() == {'type': 'connected', 'data': []}
async def test_too_many_user_connections(app: Starlette) -> None:
async with num_active_connections(app, MAX_CONNECTIONS_PER_USER, unique_ips=False):
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(
app, f'/{uuid4()}', client_ip='127.0.0.1'
) as websocket:
await websocket.receive_json()
assert e.value.code == ERR_TOO_MANY_CONNECTIONS
async def test_invalid_request(app: Starlette) -> None:
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(app, f'/{ROOM_ID}') as client:
# Wait for the initial connection message
await client.receive_json()
await client.send('{invalid_json')
await client.receive_json()
assert e.value.code == ERR_INVALID_REQUEST
async def test_too_many_room_connections(app: Starlette) -> None:
async with num_active_connections(
app, MAX_CONNECTIONS_PER_ROOM, unique_rooms=False
):
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(
app, f'/{ROOM_ID}', client_ip='127.0.0.255'
) as client:
await client.receive_json()
assert e.value.code == ERR_ROOM_FULL
async def test_bypass_max_connections_rate_limit_invalid_key(app: Starlette) -> None:
headers = {BYPASS_RATE_LIMIT_HEADER: "invalid-key"}
async with num_active_connections(
app, MAX_CONNECTIONS_PER_USER, unique_ips=False, headers=headers
):
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(
app, f'/{uuid4()}', headers=headers
) as websocket:
await websocket.receive_json()
assert e.value.code == ERR_TOO_MANY_CONNECTIONS
async def test_add_token(app: Starlette) -> None:
async with emulated_client.connect(app, f'/{ROOM_ID}') as client:
# Grab the connection message
await client.receive_json()
await client.send_json(
{
'request_id': TEST_REQUEST_ID,
'actions': [TEST_UPSERT_TOKEN],
}
)
assert_matches(
await client.receive_json(),
{
'type': 'update',
'request_id': TEST_REQUEST_ID,
'actions': [TEST_UPSERT_TOKEN],
},
)
async def test_invalid_uuid(app: Starlette) -> None:
with pytest.raises(WebsocketClosed) as e:
async with emulated_client.connect(app, '/invalid-uuid'):
pass
assert e.value.code == ERR_INVALID_UUID
async def test_connect(app: Starlette) -> None:
async with emulated_client.connect(app, f'/{ROOM_ID}') as client:
assert await client.receive_json() == {'type': 'connected', 'data': []}
