rechain
5 months ago
4 changed files with 626 additions and 0 deletions
@ -0,0 +1,106 @@ |
|||
from collections.abc import Sequence |
|||
from typing import Any |
|||
|
|||
from fastapi import routing |
|||
from fastapi.encoders import jsonable_encoder |
|||
from starlette.routing import BaseRoute |
|||
|
|||
|
|||
def get_asyncapi_channel( |
|||
*, |
|||
route: routing.APIWebSocketRoute, |
|||
) -> dict[str, Any]: |
|||
"""Generate AsyncAPI channel definition for a WebSocket route.""" |
|||
channel: dict[str, Any] = {} |
|||
|
|||
# WebSocket channels typically have subscribe operation |
|||
# (client subscribes to receive messages from server) |
|||
operation: dict[str, Any] = { |
|||
"operationId": route.name or f"websocket_{route.path_format}", |
|||
} |
|||
|
|||
# Basic message schema - can be enhanced later with actual message types |
|||
# For WebSockets, messages can be sent in both directions |
|||
message: dict[str, Any] = { |
|||
"contentType": "application/json", |
|||
} |
|||
|
|||
operation["message"] = message |
|||
channel["subscribe"] = operation |
|||
|
|||
# WebSockets are bidirectional, so we also include publish |
|||
# (client can publish messages to server) |
|||
publish_operation: dict[str, Any] = { |
|||
"operationId": f"{route.name or f'websocket_{route.path_format}'}_publish", |
|||
"message": message, |
|||
} |
|||
channel["publish"] = publish_operation |
|||
|
|||
return channel |
|||
|
|||
|
|||
def get_asyncapi( |
|||
*, |
|||
title: str, |
|||
version: str, |
|||
asyncapi_version: str = "2.6.0", |
|||
summary: str | None = None, |
|||
description: str | None = None, |
|||
routes: Sequence[BaseRoute], |
|||
servers: list[dict[str, str | Any]] | None = None, |
|||
terms_of_service: str | None = None, |
|||
contact: dict[str, str | Any] | None = None, |
|||
license_info: dict[str, str | Any] | None = None, |
|||
external_docs: dict[str, Any] | None = None, |
|||
) -> dict[str, Any]: |
|||
""" |
|||
Generate AsyncAPI schema from FastAPI application routes. |
|||
|
|||
Filters for WebSocket routes and generates AsyncAPI 2.6.0 compliant schema. |
|||
""" |
|||
info: dict[str, Any] = {"title": title, "version": version} |
|||
if summary: |
|||
info["summary"] = summary |
|||
if description: |
|||
info["description"] = description |
|||
if terms_of_service: |
|||
info["termsOfService"] = terms_of_service |
|||
if contact: |
|||
info["contact"] = contact |
|||
if license_info: |
|||
info["license"] = license_info |
|||
|
|||
output: dict[str, Any] = {"asyncapi": asyncapi_version, "info": info} |
|||
|
|||
# Add default WebSocket server if no servers provided and we have WebSocket routes |
|||
websocket_routes = [ |
|||
route for route in routes or [] if isinstance(route, routing.APIWebSocketRoute) |
|||
] |
|||
if websocket_routes and not servers: |
|||
# Default WebSocket server - can be overridden by providing servers parameter |
|||
output["servers"] = [ |
|||
{ |
|||
"url": "ws://localhost:8000", |
|||
"protocol": "ws", |
|||
"description": "WebSocket server", |
|||
} |
|||
] |
|||
elif servers: |
|||
output["servers"] = servers |
|||
|
|||
channels: dict[str, dict[str, Any]] = {} |
|||
|
|||
# Filter routes to only include WebSocket routes |
|||
for route in routes or []: |
|||
if isinstance(route, routing.APIWebSocketRoute): |
|||
channel = get_asyncapi_channel(route=route) |
|||
if channel: |
|||
channels[route.path_format] = channel |
|||
|
|||
output["channels"] = channels |
|||
|
|||
if external_docs: |
|||
output["externalDocs"] = external_docs |
|||
|
|||
return jsonable_encoder(output, by_alias=True, exclude_none=True) |
|||
|
|||
@ -0,0 +1,262 @@ |
|||
from fastapi import FastAPI, WebSocket |
|||
from fastapi.testclient import TestClient |
|||
|
|||
|
|||
def test_asyncapi_schema(): |
|||
"""Test AsyncAPI schema endpoint with WebSocket routes.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
@app.websocket("/ws/{item_id}") |
|||
async def websocket_with_param(websocket: WebSocket, item_id: str): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 200, response.text |
|||
schema = response.json() |
|||
assert schema["asyncapi"] == "2.6.0" |
|||
assert schema["info"]["title"] == "Test API" |
|||
assert schema["info"]["version"] == "1.0.0" |
|||
assert "channels" in schema |
|||
assert "/ws" in schema["channels"] |
|||
assert "/ws/{item_id}" in schema["channels"] |
|||
|
|||
|
|||
def test_asyncapi_no_websockets(): |
|||
"""Test AsyncAPI schema with no WebSocket routes.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.get("/") |
|||
def read_root(): |
|||
return {"message": "Hello World"} |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 200, response.text |
|||
schema = response.json() |
|||
assert schema["asyncapi"] == "2.6.0" |
|||
assert schema["info"]["title"] == "Test API" |
|||
assert schema["channels"] == {} |
|||
|
|||
|
|||
def test_asyncapi_caching(): |
|||
"""Test that AsyncAPI schema is cached.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
schema1 = app.asyncapi() |
|||
schema2 = app.asyncapi() |
|||
# Should return the same object (identity check) |
|||
assert schema1 is schema2 |
|||
|
|||
|
|||
def test_asyncapi_ui(): |
|||
"""Test AsyncAPI UI endpoint.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi-docs") |
|||
assert response.status_code == 200, response.text |
|||
assert response.headers["content-type"] == "text/html; charset=utf-8" |
|||
assert "@asyncapi/react-component" in response.text |
|||
assert "/asyncapi.json" in response.text |
|||
|
|||
|
|||
def test_asyncapi_ui_navigation(): |
|||
"""Test navigation links in AsyncAPI UI.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi-docs") |
|||
assert response.status_code == 200, response.text |
|||
# Should contain link to OpenAPI docs |
|||
assert "/docs" in response.text |
|||
assert "OpenAPI Docs" in response.text |
|||
|
|||
|
|||
def test_swagger_ui_asyncapi_navigation(): |
|||
"""Test navigation link to AsyncAPI in Swagger UI.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.get("/") |
|||
def read_root(): |
|||
return {"message": "Hello World"} |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/docs") |
|||
assert response.status_code == 200, response.text |
|||
# Should contain link to AsyncAPI docs |
|||
assert "/asyncapi-docs" in response.text |
|||
assert "AsyncAPI Docs" in response.text |
|||
|
|||
|
|||
def test_asyncapi_custom_urls(): |
|||
"""Test custom AsyncAPI URLs.""" |
|||
app = FastAPI( |
|||
title="Test API", |
|||
version="1.0.0", |
|||
asyncapi_url="/custom/asyncapi.json", |
|||
asyncapi_docs_url="/custom/asyncapi-docs", |
|||
) |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
# Test custom JSON endpoint |
|||
response = client.get("/custom/asyncapi.json") |
|||
assert response.status_code == 200, response.text |
|||
schema = response.json() |
|||
assert schema["asyncapi"] == "2.6.0" |
|||
|
|||
# Test custom UI endpoint |
|||
response = client.get("/custom/asyncapi-docs") |
|||
assert response.status_code == 200, response.text |
|||
assert "/custom/asyncapi.json" in response.text |
|||
|
|||
# Default endpoints should not exist |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 404 |
|||
response = client.get("/asyncapi-docs") |
|||
assert response.status_code == 404 |
|||
|
|||
|
|||
def test_asyncapi_disabled(): |
|||
"""Test when AsyncAPI is disabled.""" |
|||
app = FastAPI( |
|||
title="Test API", |
|||
version="1.0.0", |
|||
asyncapi_url=None, |
|||
) |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
# Endpoints should return 404 |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 404 |
|||
response = client.get("/asyncapi-docs") |
|||
assert response.status_code == 404 |
|||
|
|||
|
|||
def test_asyncapi_channel_structure(): |
|||
"""Test AsyncAPI channel structure.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 200, response.text |
|||
schema = response.json() |
|||
channel = schema["channels"]["/ws"] |
|||
assert "subscribe" in channel |
|||
assert "operationId" in channel["subscribe"] |
|||
assert "message" in channel["subscribe"] |
|||
|
|||
|
|||
def test_asyncapi_multiple_websockets(): |
|||
"""Test AsyncAPI with multiple WebSocket routes.""" |
|||
app = FastAPI(title="Test API", version="1.0.0") |
|||
|
|||
@app.websocket("/ws1") |
|||
async def websocket1(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
@app.websocket("/ws2") |
|||
async def websocket2(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
@app.websocket("/ws3/{param}") |
|||
async def websocket3(websocket: WebSocket, param: str): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 200, response.text |
|||
schema = response.json() |
|||
assert len(schema["channels"]) == 3 |
|||
assert "/ws1" in schema["channels"] |
|||
assert "/ws2" in schema["channels"] |
|||
assert "/ws3/{param}" in schema["channels"] |
|||
|
|||
|
|||
def test_asyncapi_with_metadata(): |
|||
"""Test AsyncAPI schema includes app metadata.""" |
|||
app = FastAPI( |
|||
title="My API", |
|||
version="2.0.0", |
|||
summary="Test summary", |
|||
description="Test description", |
|||
) |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi.json") |
|||
assert response.status_code == 200, response.text |
|||
schema = response.json() |
|||
assert schema["info"]["title"] == "My API" |
|||
assert schema["info"]["version"] == "2.0.0" |
|||
assert schema["info"]["summary"] == "Test summary" |
|||
assert schema["info"]["description"] == "Test description" |
|||
|
|||
|
|||
def test_asyncapi_ui_no_docs_url(): |
|||
"""Test AsyncAPI UI when docs_url is None.""" |
|||
app = FastAPI( |
|||
title="Test API", |
|||
version="1.0.0", |
|||
docs_url=None, |
|||
) |
|||
|
|||
@app.websocket("/ws") |
|||
async def websocket_endpoint(websocket: WebSocket): |
|||
await websocket.accept() |
|||
await websocket.close() |
|||
|
|||
client = TestClient(app) |
|||
response = client.get("/asyncapi-docs") |
|||
assert response.status_code == 200, response.text |
|||
# Should not contain link to /docs if docs_url is None |
|||
# But navigation should still work (just won't show the link) |
|||
assert "/asyncapi.json" in response.text |
|||
Loading…
Reference in new issue