#![allow(clippy::expect_used)]
use std::sync::{Arc, Mutex, Once};
use tensor_wasm_api::server::__reset_host_validation_warn_for_test;
use tensor_wasm_api::{build_router, AppState, ENV_API_TOKENS, ENV_TRUSTED_HOSTS};
use tracing::Subscriber;
use tracing_subscriber::layer::{Context, SubscriberExt};
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::Layer;
const T16_WARN_TARGET: &str = "tensor_wasm_api::server";
const T16_MESSAGE_FRAGMENT: &str = "TENSOR_WASM_API_TRUSTED_HOSTS";
#[derive(Debug, Clone)]
struct CapturedEvent {
target: String,
message: String,
}
#[derive(Clone)]
struct CapturingLayer {
events: Arc<Mutex<Vec<CapturedEvent>>>,
}
impl<S> Layer<S> for CapturingLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
if *event.metadata().level() != tracing::Level::WARN {
return;
}
let target = event.metadata().target().to_string();
let mut visitor = MessageVisitor::default();
event.record(&mut visitor);
let captured = CapturedEvent {
target,
message: visitor.message,
};
if let Ok(mut guard) = self.events.lock() {
guard.push(captured);
}
}
}
#[derive(Default)]
struct MessageVisitor {
message: String,
}
impl tracing::field::Visit for MessageVisitor {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
use std::fmt::Write;
if field.name() == "message" {
let _ = write!(self.message, "{value:?}");
} else {
let _ = write!(self.message, " {}={value:?}", field.name());
}
}
}
fn install_capturing_subscriber() -> Arc<Mutex<Vec<CapturedEvent>>> {
static INIT: Once = Once::new();
static EVENTS: std::sync::OnceLock<Arc<Mutex<Vec<CapturedEvent>>>> = std::sync::OnceLock::new();
let events = EVENTS
.get_or_init(|| Arc::new(Mutex::new(Vec::new())))
.clone();
INIT.call_once(|| {
let layer = CapturingLayer {
events: Arc::clone(&events),
};
let _ = tracing_subscriber::registry().with(layer).try_init();
});
events
}
fn drain_new_events(
events: &Arc<Mutex<Vec<CapturedEvent>>>,
baseline: usize,
) -> Vec<CapturedEvent> {
let guard = events.lock().expect("capture mutex poisoned");
guard[baseline..].to_vec()
}
fn t16_warn_present(events: &[CapturedEvent]) -> bool {
events
.iter()
.any(|e| e.target == T16_WARN_TARGET && e.message.contains(T16_MESSAGE_FRAGMENT))
}
#[test]
fn warns_when_tokens_set_but_trusted_hosts_unset() {
let events = install_capturing_subscriber();
temp_env::with_vars(
[
(ENV_API_TOKENS, Some("test")),
(ENV_TRUSTED_HOSTS, None::<&str>),
],
|| {
__reset_host_validation_warn_for_test();
let baseline = events.lock().expect("baseline lock").len();
let state = Arc::new(AppState::default());
let _router = build_router(state);
let new_events = drain_new_events(&events, baseline);
assert!(
t16_warn_present(&new_events),
"expected T16 warn against target {T16_WARN_TARGET:?} containing \
{T16_MESSAGE_FRAGMENT:?}, got: {new_events:#?}",
);
},
);
}
#[test]
fn no_warn_when_both_vars_set() {
let events = install_capturing_subscriber();
temp_env::with_vars(
[
(ENV_API_TOKENS, Some("test")),
(ENV_TRUSTED_HOSTS, Some("api.example.com")),
],
|| {
__reset_host_validation_warn_for_test();
let baseline = events.lock().expect("baseline lock").len();
let state = Arc::new(AppState::default());
let _router = build_router(state);
let new_events = drain_new_events(&events, baseline);
assert!(
!t16_warn_present(&new_events),
"T16 warn must NOT fire when TENSOR_WASM_API_TRUSTED_HOSTS \
is set; got: {new_events:#?}",
);
},
);
}
#[test]
fn no_warn_when_neither_var_set() {
let events = install_capturing_subscriber();
temp_env::with_vars(
[
(ENV_API_TOKENS, None::<&str>),
(ENV_TRUSTED_HOSTS, None::<&str>),
],
|| {
__reset_host_validation_warn_for_test();
let baseline = events.lock().expect("baseline lock").len();
let state = Arc::new(AppState::default());
let _router = build_router(state);
let new_events = drain_new_events(&events, baseline);
assert!(
!t16_warn_present(&new_events),
"T16 warn must NOT fire in pure dev mode (neither env var \
set); got: {new_events:#?}",
);
},
);
}