mirror of https://github.com/encode/starlette.git
147 lines
3.8 KiB
Python
147 lines
3.8 KiB
Python
from starlette.responses import Response
|
|
from starlette.testclient import TestClient
|
|
from starlette.exceptions import ExceptionMiddleware
|
|
from starlette.routing import Route, Mount, NoMatchFound, Router, WebSocketRoute
|
|
from starlette.websockets import WebSocket, WebSocketDisconnect
|
|
import pytest
|
|
|
|
|
|
def homepage(request):
|
|
return Response("Hello, world", media_type="text/plain")
|
|
|
|
|
|
def users(request):
|
|
return Response("All users", media_type="text/plain")
|
|
|
|
|
|
def user(request):
|
|
content = "User " + request.path_params["username"]
|
|
return Response(content, media_type="text/plain")
|
|
|
|
|
|
def staticfiles(request):
|
|
return Response("xxxxx", media_type="image/png")
|
|
|
|
|
|
app = Router(
|
|
[
|
|
Route("/", endpoint=homepage, methods=["GET"]),
|
|
Mount(
|
|
"/users",
|
|
app=Router(
|
|
[Route("", endpoint=users), Route("/{username}", endpoint=user)]
|
|
),
|
|
),
|
|
Mount("/static", app=staticfiles),
|
|
]
|
|
)
|
|
|
|
|
|
@app.route("/func")
|
|
def func_homepage(request):
|
|
return Response("Hello, world!", media_type="text/plain")
|
|
|
|
|
|
@app.websocket_route("/ws")
|
|
async def websocket_endpoint(session):
|
|
await session.accept()
|
|
await session.send_text("Hello, world!")
|
|
await session.close()
|
|
|
|
|
|
@app.websocket_route("/ws/{room}")
|
|
async def websocket_params(session):
|
|
await session.accept()
|
|
await session.send_text("Hello, %s!" % session.path_params["room"])
|
|
await session.close()
|
|
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
def test_router():
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert response.text == "Hello, world"
|
|
|
|
response = client.post("/")
|
|
assert response.status_code == 405
|
|
assert response.text == "Method Not Allowed"
|
|
|
|
response = client.get("/foo")
|
|
assert response.status_code == 404
|
|
assert response.text == "Not Found"
|
|
|
|
response = client.get("/users")
|
|
assert response.status_code == 200
|
|
assert response.text == "All users"
|
|
|
|
response = client.get("/users/tomchristie")
|
|
assert response.status_code == 200
|
|
assert response.text == "User tomchristie"
|
|
|
|
response = client.get("/static/123")
|
|
assert response.status_code == 200
|
|
assert response.text == "xxxxx"
|
|
|
|
|
|
def test_url_for():
|
|
assert app.url_for("homepage") == "/"
|
|
assert app.url_for("user", username="tomchristie") == "/users/tomchristie"
|
|
assert app.url_for("websocket_endpoint") == "/ws"
|
|
with pytest.raises(NoMatchFound):
|
|
assert app.url_for("broken")
|
|
|
|
|
|
def test_router_add_route():
|
|
response = client.get("/func")
|
|
assert response.status_code == 200
|
|
assert response.text == "Hello, world!"
|
|
|
|
|
|
def test_router_add_websocket_route():
|
|
with client.websocket_connect("/ws") as session:
|
|
text = session.receive_text()
|
|
assert text == "Hello, world!"
|
|
|
|
with client.websocket_connect("/ws/test") as session:
|
|
text = session.receive_text()
|
|
assert text == "Hello, test!"
|
|
|
|
|
|
def http_endpoint(request):
|
|
return Response("Hello, world", media_type="text/plain")
|
|
|
|
|
|
class WebsocketEndpoint:
|
|
def __init__(self, scope):
|
|
self.scope = scope
|
|
|
|
async def __call__(self, receive, send):
|
|
session = WebSocket(scope=self.scope, receive=receive, send=send)
|
|
await session.accept()
|
|
await session.send_json({"hello": "world"})
|
|
await session.close()
|
|
|
|
|
|
mixed_protocol_app = Router(
|
|
routes=[
|
|
Route("/", endpoint=http_endpoint),
|
|
WebSocketRoute("/", endpoint=WebsocketEndpoint),
|
|
]
|
|
)
|
|
|
|
|
|
def test_protocol_switch():
|
|
client = TestClient(mixed_protocol_app)
|
|
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert response.text == "Hello, world"
|
|
|
|
with client.websocket_connect("/") as session:
|
|
assert session.receive_json() == {"hello": "world"}
|
|
|
|
with pytest.raises(WebSocketDisconnect):
|
|
client.websocket_connect("/404")
|