mirror of https://github.com/tiangolo/fastapi.git
47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import pytest
|
|
from fastapi import FastAPI, Security
|
|
from fastapi.security import HTTPAuthorizationCredentials, HTTPDigest
|
|
from fastapi.testclient import TestClient
|
|
from starlette.testclient import WebSocketDenialResponse
|
|
from starlette.websockets import WebSocket
|
|
|
|
app = FastAPI()
|
|
|
|
security = HTTPDigest()
|
|
|
|
|
|
@app.websocket("/ws/users/me")
|
|
async def read_current_user(
|
|
websocket: WebSocket,
|
|
credentials: HTTPAuthorizationCredentials = Security(security),
|
|
):
|
|
await websocket.accept()
|
|
await websocket.send_json(
|
|
{"scheme": credentials.scheme, "credentials": credentials.credentials}
|
|
)
|
|
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
def test_security_http_digest_ws():
|
|
with client.websocket_connect(
|
|
"/ws/users/me", headers={"Authorization": "Digest foobar"}
|
|
) as websocket:
|
|
data = websocket.receive_json()
|
|
assert data == {"scheme": "Digest", "credentials": "foobar"}
|
|
|
|
|
|
def test_security_http_digest_no_credentials_ws():
|
|
with pytest.raises(WebSocketDenialResponse):
|
|
with client.websocket_connect("/ws/users/me"):
|
|
pass
|
|
|
|
|
|
def test_security_http_digest_incorrect_scheme_ws():
|
|
with pytest.raises(WebSocketDenialResponse):
|
|
with client.websocket_connect(
|
|
"/ws/users/me", headers={"Authorization": "Basic notreally"}
|
|
):
|
|
pass
|