mirror of https://github.com/tiangolo/fastapi.git
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
import pytest
|
|
from fastapi import Depends, FastAPI, Security
|
|
from fastapi.security import APIKeyQuery
|
|
from fastapi.testclient import TestClient
|
|
from pydantic import BaseModel
|
|
from starlette.testclient import WebSocketDenialResponse
|
|
from starlette.websockets import WebSocket
|
|
|
|
app = FastAPI()
|
|
|
|
api_key = APIKeyQuery(name="key")
|
|
|
|
|
|
class User(BaseModel):
|
|
username: str
|
|
|
|
|
|
def get_current_user(oauth_header: str = Security(api_key)):
|
|
user = User(username=oauth_header)
|
|
return user
|
|
|
|
|
|
@app.websocket("/ws/users/me")
|
|
async def read_current_user(
|
|
websocket: WebSocket, current_user: User = Depends(get_current_user)
|
|
):
|
|
await websocket.accept()
|
|
await websocket.send_text(current_user.username)
|
|
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
def test_security_api_key_query_ws():
|
|
with client.websocket_connect("/ws/users/me?key=secret") as websocket:
|
|
data = websocket.receive_text()
|
|
assert data == "secret"
|
|
|
|
|
|
def test_security_api_key_query_no_key_ws():
|
|
with pytest.raises(WebSocketDenialResponse):
|
|
with client.websocket_connect("/ws/users/me"):
|
|
pass
|