use axum::{
routing::{delete, get, post},
Router,
};
use tower::ServiceBuilder;
use tower_http::{
compression::CompressionLayer,
cors::CorsLayer,
trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
};
use tracing::Level;
use super::{
handlers::{
convert_document, delete_pdf_job, get_health, get_ocr_results, get_ocr_usage,
get_pdf_status, process_latex, process_pdf, process_strokes, process_text,
stream_pdf_results,
},
middleware::{auth_middleware, rate_limit_middleware},
state::AppState,
};
pub fn router(state: AppState) -> Router {
let api_routes = Router::new()
.route("/v3/text", post(process_text))
.route("/v3/strokes", post(process_strokes))
.route("/v3/latex", post(process_latex))
.route("/v3/pdf", post(process_pdf))
.route("/v3/pdf/:id", get(get_pdf_status))
.route("/v3/pdf/:id", delete(delete_pdf_job))
.route("/v3/pdf/:id/stream", get(stream_pdf_results))
.route("/v3/converter", post(convert_document))
.route("/v3/ocr-results", get(get_ocr_results))
.route("/v3/ocr-usage", get(get_ocr_usage))
.layer(
ServiceBuilder::new()
.layer(axum::middleware::from_fn_with_state(
state.clone(),
auth_middleware,
))
.layer(axum::middleware::from_fn_with_state(
state.clone(),
rate_limit_middleware,
)),
);
let health_routes = Router::new().route("/health", get(get_health));
Router::new()
.merge(api_routes)
.merge(health_routes)
.layer(
ServiceBuilder::new()
.layer(
TraceLayer::new_for_http()
.make_span_with(DefaultMakeSpan::new().level(Level::INFO))
.on_response(DefaultOnResponse::new().level(Level::INFO)),
)
.layer(CorsLayer::permissive())
.layer(CompressionLayer::new()),
)
.with_state(state)
}
#[cfg(test)]
mod tests {
use super::*;
use axum::body::Body;
use axum::http::{Request, StatusCode};
use tower::ServiceExt;
#[tokio::test]
async fn test_health_endpoint() {
let state = AppState::new();
let app = router(state);
let response = app
.oneshot(
Request::builder()
.uri("/health")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
}