use axum::{
extract::{Query, State},
http::StatusCode,
response::Json,
};
use mockforge_core::ab_testing::analytics::ABTestReport;
use mockforge_core::ab_testing::{
ABTestConfig, VariantAnalytics, VariantComparison, VariantManager,
};
use serde::Deserialize;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{error, info};
#[derive(Clone)]
pub struct ABTestingState {
pub variant_manager: Arc<VariantManager>,
}
impl Default for ABTestingState {
fn default() -> Self {
Self::new()
}
}
impl ABTestingState {
pub fn new() -> Self {
Self {
variant_manager: Arc::new(VariantManager::new()),
}
}
}
#[derive(Debug, Deserialize)]
pub struct CreateABTestRequest {
pub test: ABTestConfig,
}
#[derive(Debug, Deserialize)]
pub struct UpdateAllocationRequest {
pub allocations: Vec<mockforge_core::ab_testing::VariantAllocation>,
}
#[derive(Debug, Deserialize)]
pub struct EndpointQuery {
pub method: String,
pub path: String,
}
pub async fn create_ab_test(
State(state): State<ABTestingState>,
Json(req): Json<CreateABTestRequest>,
) -> Result<Json<ABTestConfig>, StatusCode> {
info!("Creating A/B test: {}", req.test.test_name);
if let Err(e) = req.test.validate_allocations() {
error!("Invalid A/B test configuration: {}", e);
return Err(StatusCode::BAD_REQUEST);
}
match state.variant_manager.register_test(req.test.clone()).await {
Ok(_) => {
info!("A/B test created successfully: {}", req.test.test_name);
Ok(Json(req.test))
}
Err(e) => {
error!("Failed to create A/B test: {}", e);
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
}
pub async fn get_ab_test(
State(state): State<ABTestingState>,
Query(params): Query<EndpointQuery>,
) -> Result<Json<ABTestConfig>, StatusCode> {
state
.variant_manager
.get_test(¶ms.method, ¶ms.path)
.await
.map(Json)
.ok_or(StatusCode::NOT_FOUND)
}
pub async fn list_ab_tests(
State(state): State<ABTestingState>,
) -> Result<Json<Vec<ABTestConfig>>, StatusCode> {
let tests = state.variant_manager.list_tests().await;
Ok(Json(tests))
}
pub async fn delete_ab_test(
State(state): State<ABTestingState>,
Query(params): Query<EndpointQuery>,
) -> Result<Json<Value>, StatusCode> {
match state.variant_manager.remove_test(¶ms.method, ¶ms.path).await {
Ok(_) => {
info!("A/B test deleted: {} {}", params.method, params.path);
Ok(Json(serde_json::json!({
"success": true,
"message": format!("A/B test deleted for {} {}", params.method, params.path)
})))
}
Err(e) => {
error!("Failed to delete A/B test: {}", e);
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
}
pub async fn get_ab_test_analytics(
State(state): State<ABTestingState>,
Query(params): Query<EndpointQuery>,
) -> Result<Json<ABTestReport>, StatusCode> {
let test_config = state
.variant_manager
.get_test(¶ms.method, ¶ms.path)
.await
.ok_or(StatusCode::NOT_FOUND)?;
let variant_analytics =
state.variant_manager.get_endpoint_analytics(¶ms.method, ¶ms.path).await;
let report = ABTestReport::new(test_config, variant_analytics);
Ok(Json(report))
}
pub async fn get_variant_analytics(
State(state): State<ABTestingState>,
Query(params): Query<HashMap<String, String>>,
) -> Result<Json<VariantAnalytics>, StatusCode> {
let method = params.get("method").ok_or(StatusCode::BAD_REQUEST)?;
let path = params.get("path").ok_or(StatusCode::BAD_REQUEST)?;
let variant_id = params.get("variant_id").ok_or(StatusCode::BAD_REQUEST)?;
state
.variant_manager
.get_variant_analytics(method, path, variant_id)
.await
.map(Json)
.ok_or(StatusCode::NOT_FOUND)
}
pub async fn compare_variants(
State(state): State<ABTestingState>,
Query(params): Query<HashMap<String, String>>,
) -> Result<Json<VariantComparison>, StatusCode> {
let method = params.get("method").ok_or(StatusCode::BAD_REQUEST)?;
let path = params.get("path").ok_or(StatusCode::BAD_REQUEST)?;
let variant_a_id = params.get("variant_a").ok_or(StatusCode::BAD_REQUEST)?;
let variant_b_id = params.get("variant_b").ok_or(StatusCode::BAD_REQUEST)?;
let analytics_a = state
.variant_manager
.get_variant_analytics(method, path, variant_a_id)
.await
.ok_or(StatusCode::NOT_FOUND)?;
let analytics_b = state
.variant_manager
.get_variant_analytics(method, path, variant_b_id)
.await
.ok_or(StatusCode::NOT_FOUND)?;
let comparison = VariantComparison::new(&analytics_a, &analytics_b);
Ok(Json(comparison))
}
pub async fn update_allocations(
State(state): State<ABTestingState>,
Query(params): Query<EndpointQuery>,
Json(req): Json<UpdateAllocationRequest>,
) -> Result<Json<ABTestConfig>, StatusCode> {
let mut test_config = state
.variant_manager
.get_test(¶ms.method, ¶ms.path)
.await
.ok_or(StatusCode::NOT_FOUND)?;
test_config.allocations = req.allocations;
if let Err(e) = test_config.validate_allocations() {
error!("Invalid allocations: {}", e);
return Err(StatusCode::BAD_REQUEST);
}
match state.variant_manager.register_test(test_config.clone()).await {
Ok(_) => {
info!("Updated allocations for {} {}", params.method, params.path);
Ok(Json(test_config))
}
Err(e) => {
error!("Failed to update allocations: {}", e);
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
}
pub async fn toggle_ab_test(
State(state): State<ABTestingState>,
Query(params): Query<HashMap<String, String>>,
) -> Result<Json<ABTestConfig>, StatusCode> {
let method = params.get("method").ok_or(StatusCode::BAD_REQUEST)?;
let path = params.get("path").ok_or(StatusCode::BAD_REQUEST)?;
let enabled = params
.get("enabled")
.and_then(|v| v.parse::<bool>().ok())
.ok_or(StatusCode::BAD_REQUEST)?;
let mut test_config = state
.variant_manager
.get_test(method, path)
.await
.ok_or(StatusCode::NOT_FOUND)?;
test_config.enabled = enabled;
match state.variant_manager.register_test(test_config.clone()).await {
Ok(_) => {
info!(
"{} A/B test for {} {}",
if enabled { "Enabled" } else { "Disabled" },
method,
path
);
Ok(Json(test_config))
}
Err(e) => {
error!("Failed to toggle A/B test: {}", e);
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
}
pub fn ab_testing_router(state: ABTestingState) -> axum::Router {
use axum::routing::{delete, get, patch, post, put};
use axum::Router;
Router::new()
.route("/api/v1/ab-tests", post(create_ab_test).get(list_ab_tests))
.route("/api/v1/ab-tests/analytics", get(get_ab_test_analytics))
.route("/api/v1/ab-tests/variants/analytics", get(get_variant_analytics))
.route("/api/v1/ab-tests/variants/compare", get(compare_variants))
.route("/api/v1/ab-tests/allocations", put(update_allocations))
.route("/api/v1/ab-tests/enable", patch(toggle_ab_test))
.route("/api/v1/ab-tests/delete", delete(delete_ab_test))
.with_state(state)
}