mirror of https://github.com/tiangolo/fastapi.git
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import pytest
|
|
from fastapi import FastAPI, Security
|
|
from fastapi.security import OAuth2AuthorizationCodeBearer
|
|
from fastapi.testclient import TestClient
|
|
from starlette.testclient import WebSocketDenialResponse
|
|
from starlette.websockets import WebSocket
|
|
|
|
app = FastAPI()
|
|
|
|
oauth2_scheme = OAuth2AuthorizationCodeBearer(
|
|
authorizationUrl="/api/oauth/authorize",
|
|
tokenUrl="/api/oauth/token",
|
|
scopes={"read": "Read access", "write": "Write access"},
|
|
)
|
|
|
|
|
|
@app.websocket("/ws/admin")
|
|
async def read_admin(websocket: WebSocket, token: str = Security(oauth2_scheme)):
|
|
await websocket.accept()
|
|
await websocket.send_text(token)
|
|
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
def test_security_oauth2_authorization_code_bearer_ws():
|
|
with client.websocket_connect(
|
|
"/ws/admin", headers={"Authorization": "Bearer faketoken"}
|
|
) as websocket:
|
|
data = websocket.receive_text()
|
|
assert data == "faketoken"
|
|
|
|
|
|
def test_security_oauth2_authorization_code_bearer_no_header_ws():
|
|
with pytest.raises(WebSocketDenialResponse):
|
|
with client.websocket_connect("/ws/admin"):
|
|
pass
|
|
|
|
|
|
def test_security_oauth2_authorization_code_bearer_wrong_scheme_ws():
|
|
with pytest.raises(WebSocketDenialResponse):
|
|
with client.websocket_connect(
|
|
"/ws/admin", headers={"Authorization": "Basic nope"}
|
|
):
|
|
pass
|