feat: improve test reliability and add request tracing

- Add retry configuration for flaky tests (2 retries for default, 3 for
OAuth)
- Configure test groups with proper concurrency limits (serial: 1,
integration: 4)
- Add tower-http tracing layer with custom span formatting for HTTP
requests
- Simplify database pool handling by removing unnecessary Arc wrapper
- Improve test context setup with better logging and error handling
- Refactor user creation parameters for better clarity and consistency
- Add debug logging for OAuth cookie handling
This commit is contained in:
Ryan Walters
2025-09-19 17:35:53 -05:00
parent e1b266f3b2
commit 0b5aeceb51
10 changed files with 446 additions and 125 deletions
+34 -4
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{Notify, RwLock};
use tokio::task::JoinHandle;
use tracing::info_span;
use crate::data::pool::PgPool;
use crate::{auth::AuthRegistry, config::Config, image::ImageStorage, routes};
@@ -36,7 +37,7 @@ pub struct AppState {
pub sessions: Arc<DashMap<String, crate::auth::provider::AuthUser>>,
pub jwt_encoding_key: Arc<EncodingKey>,
pub jwt_decoding_key: Arc<DecodingKey>,
pub db: Arc<PgPool>,
pub db: PgPool,
pub health: Arc<RwLock<Health>>,
pub image_storage: Arc<ImageStorage>,
pub healthchecker_task: Arc<RwLock<Option<JoinHandle<()>>>>,
@@ -71,7 +72,7 @@ impl AppState {
sessions: Arc::new(DashMap::new()),
jwt_encoding_key: Arc::new(EncodingKey::from_secret(jwt_secret.as_bytes())),
jwt_decoding_key: Arc::new(DecodingKey::from_secret(jwt_secret.as_bytes())),
db: Arc::new(db),
db: db,
health: Arc::new(RwLock::new(Health::default())),
image_storage,
healthchecker_task: Arc::new(RwLock::new(None)),
@@ -100,7 +101,7 @@ impl AppState {
}
// Run the actual health check
let ok = sqlx::query("SELECT 1").execute(&*db_pool).await.is_ok();
let ok = sqlx::query("SELECT 1").execute(&db_pool).await.is_ok();
{
let mut h = health_state.write().await;
h.set_database(ok);
@@ -127,13 +128,35 @@ impl AppState {
/// Force an immediate health check (debug mode only)
pub async fn check_health(&self) -> bool {
let ok = sqlx::query("SELECT 1").execute(&*self.db).await.is_ok();
let ok = sqlx::query("SELECT 1").execute(&self.db).await.is_ok();
let mut h = self.health.write().await;
h.set_database(ok);
ok
}
}
/// Create a custom span for HTTP requests with reduced verbosity
pub fn make_span<B>(request: &axum::http::Request<B>) -> tracing::Span {
let path = request
.uri()
.path_and_query()
.map(|v| v.as_str())
.unwrap_or_else(|| request.uri().path());
if request.method() == axum::http::Method::GET {
info_span!(
"request",
path = %path,
)
} else {
info_span!(
"request",
method = %request.method(),
path = %path,
)
}
}
/// Create the application router with all routes and middleware
pub fn create_router(app_state: AppState) -> Router {
Router::new()
@@ -147,6 +170,13 @@ pub fn create_router(app_state: AppState) -> Router {
.with_state(app_state)
.layer(CookieLayer::default())
.layer(axum::middleware::from_fn(inject_server_header))
.layer(
tower_http::trace::TraceLayer::new_for_http()
.make_span_with(make_span)
.on_request(|_request: &axum::http::Request<axum::body::Body>, _span: &tracing::Span| {
// Disable request logging by doing nothing
}),
)
}
/// Inject the server header into responses
+10 -10
View File
@@ -68,10 +68,10 @@ pub async fn link_oauth_account(
pub async fn create_user(
pool: &sqlx::PgPool,
username: &str,
display_name: Option<&str>,
email: Option<&str>,
avatar_url: Option<&str>,
provider_username: &str,
provider_display_name: Option<&str>,
provider_email: Option<&str>,
provider_avatar_url: Option<&str>,
provider: &str,
provider_user_id: &str,
) -> Result<User, sqlx::Error> {
@@ -82,20 +82,20 @@ pub async fn create_user(
RETURNING id, email, created_at, updated_at
"#,
)
.bind(email)
.bind(provider_email)
.fetch_one(pool)
.await?;
// Create oauth link
let _ = link_oauth_account(
let _linked = link_oauth_account(
pool,
user.id,
provider,
provider_user_id,
email,
Some(username),
display_name,
avatar_url,
provider_email,
Some(provider_username),
provider_display_name,
provider_avatar_url,
)
.await?;
+30 -28
View File
@@ -1,36 +1,38 @@
use tracing_subscriber::fmt::format::JsonFields;
use tracing_subscriber::{EnvFilter, FmtSubscriber};
use tracing_subscriber::{fmt::format::JsonFields, EnvFilter, FmtSubscriber};
use crate::config::Config;
use crate::formatter;
static SUBSCRIBER_INIT: std::sync::Once = std::sync::Once::new();
/// Configure and initialize logging for the application
pub fn setup_logging(_config: &Config) {
// Allow RUST_LOG to override levels; default to info for our crate and warn elsewhere
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new(format!("warn,{name}=info,{name}::auth=info", name = env!("CARGO_CRATE_NAME"))));
pub fn setup_logging() {
SUBSCRIBER_INIT.call_once(|| {
// Allow RUST_LOG to override levels; default to info for our crate and warn elsewhere
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new(format!("warn,{name}=info,{name}::auth=info", name = env!("CARGO_CRATE_NAME"))));
// Default to pretty for local dev; switchable later if we add CLI
let use_pretty = cfg!(debug_assertions);
// Default to pretty for local dev; switchable later if we add CLI
let use_pretty = cfg!(debug_assertions);
let subscriber: Box<dyn tracing::Subscriber + Send + Sync> = if use_pretty {
Box::new(
FmtSubscriber::builder()
.with_target(true)
.event_format(formatter::CustomPrettyFormatter)
.with_env_filter(filter)
.finish(),
)
} else {
Box::new(
FmtSubscriber::builder()
.with_target(true)
.event_format(formatter::CustomJsonFormatter)
.fmt_fields(JsonFields::new())
.with_env_filter(filter)
.finish(),
)
};
let subscriber: Box<dyn tracing::Subscriber + Send + Sync> = if use_pretty {
Box::new(
FmtSubscriber::builder()
.with_target(true)
.event_format(formatter::CustomPrettyFormatter)
.with_env_filter(filter)
.finish(),
)
} else {
Box::new(
FmtSubscriber::builder()
.with_target(true)
.event_format(formatter::CustomJsonFormatter)
.fmt_fields(JsonFields::new())
.with_env_filter(filter)
.finish(),
)
};
tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
});
}
+1 -1
View File
@@ -45,7 +45,7 @@ async fn main() {
let config: Config = config::load_config();
// Initialize tracing subscriber
logging::setup_logging(&config);
logging::setup_logging();
trace!(host = %config.host, port = config.port, shutdown_timeout_seconds = config.shutdown_timeout_seconds, "Loaded server configuration");
let addr = std::net::SocketAddr::new(config.host, config.port);
+2
View File
@@ -93,6 +93,8 @@ pub async fn oauth_callback_handler(
}
};
debug!(cookies = ?cookie.cookie().iter().collect::<Vec<_>>(), "Cookies");
// Linking or sign-in flow. Determine link intent from cookie (set at authorize time)
let link_cookie = cookie.get("link").map(|c| c.value().to_string());
if link_cookie.is_some() {