I am trying to create tests for a fast_api websocket endpoint as shown with below snippet.
import asyncio
import pytest
from sqlalchemy import (
Column,
Text
)
from gino.ext.starlette import Gino
from starlette.websockets import WebSocketDisconnect
from fastapi import Depends, FastAPI, WebSocket
from fastapi.testclient import TestClient
from fastapi import status
db = Gino(
dsn='postgresql://postgres:@0.0.0.0:5432/internal',
pool_min_size=1,
pool_max_size=16,
echo=False,
ssl=None,
use_connection_for_request=True,
retry_limit=1,
retry_interval=1,
)
app = FastAPI()
db.init_app(app)
class User(db.Model):
__tablename__ = "client_user"
__table_args__ = {"schema": 'internal'}
alias = "client_user"
username = Column(Text, nullable=False)
password = Column(Text)
async def get_current_user(websocket: WebSocket):
# for simplicity sake username is passed to sec-websocket-protocol instead of JWT
username = websocket.headers.get('sec-websocket-protocol')
# database content is
# admin/pass
user = await User.query.where(User.username == username).gino.first()
return user
@app.websocket("/websocket")
async def websocket_endpoint(
websocket: WebSocket,
user: User = Depends(get_current_user),
):
if user is None:
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
return
await websocket.accept()
print(f'User: {user.username} connected')
@pytest.fixture(scope="session")
def event_loop():
"""
This is to make the asyncio event loop shared for the whole test session, otherwise
it will be recreated for each test which will prevent using the test_db fixture.
https://github.com/FactoryBoy/factory_boy/issues/679
https://stackoverflow.com/a/56238383
"""
loop = asyncio.get_event_loop()
yield loop
loop.close()
@pytest.fixture
def client():
"""Get client for test app"""
with TestClient(app) as client:
yield client
def test_invalid_user(client):
with pytest.raises(WebSocketDisconnect):
with client.websocket_connect("/websocket", ['fake_user']):
assert True
def test_valid_user(client):
with client.websocket_connect("/websocket", ['admin']):
assert True
Description
I am trying to create tests for a fast_api websocket endpoint as shown with below snippet.
What I Did
I run test with below
here's the traceback