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