Use paused time in sampler and model catalog timeout tests (#46519)

## Why

Timeout tests can avoid waiting for real deadlines by advancing Tokio's clock, but automatic advancement can expire requests before they reach the mock server.

## What changed

Pause Tokio time in the stalled-header sampling and model catalog deadline tests. Keep polling until the mock server receives each request, then explicitly advance past the relevant timeout. Add bounded wall-clock waits and fail if a request completes before reaching the server.

Preserve assertions for sampling retry exhaustion, connection permit recovery, and the catalog's `RequestTimeout` error.

GitOrigin-RevId: 252be25d75d73c713a4b8ccac1ff20ad820df556
This commit is contained in:
jif
2026-09-18 23:58:56 +00:00
committed by copyberry
parent cdfda82386
commit 4b9e7f6473
2 changed files with 56 additions and 9 deletions
@@ -251,11 +251,32 @@ async fn stalled_http_headers_exhaust_the_sampling_retry_budget() -> Result<()>
config.provider = create_model_provider(provider, config.provider.auth_manager());
let sampler = LunaSampler::new(config);
let result = tokio::time::timeout(
tokio::time::pause();
let sample = tokio::time::timeout(
Duration::from_secs(10),
sampler.sample(sample_request("stalled-headers")),
)
.await?;
);
tokio::pin!(sample);
let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 30);
for expected_requests in 1..=3 {
// Keep the paused runtime runnable until the server records this attempt.
// Otherwise auto-advance can exhaust a retry before its request arrives.
while mock.requests().len() < expected_requests {
anyhow::ensure!(
std::time::Instant::now() < deadline,
"server did not receive request {expected_requests}"
);
tokio::select! {
result = &mut sample => {
anyhow::bail!("sampling finished before request {expected_requests}: {result:?}");
}
() = tokio::task::yield_now() => {}
}
}
tokio::time::advance(Duration::from_millis(/*millis*/ 1_001)).await;
}
let result = sample.await?;
tokio::time::resume();
assert!(matches!(
result,
Err(LunaSamplerError::Api(ApiError::Transport(
@@ -10,15 +10,15 @@ use wiremock::matchers::path;
#[tokio::test]
async fn catalog_deadline_returns_request_timeout() {
let server = MockServer::start().await;
Mock::given(method("GET"))
let mock = Mock::given(method("GET"))
.and(path("/models"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({ "models": [] }))
.set_delay(MODELS_REFRESH_TIMEOUT * 2),
.set_delay(Duration::from_secs(/*secs*/ 60)),
)
.expect(1)
.mount(&server)
.mount_as_scoped(&server)
.await;
let endpoint = OpenAiModelsEndpoint::new(
ModelProviderInfo::create_openai_provider(Some(server.uri())),
@@ -26,13 +26,39 @@ async fn catalog_deadline_returns_request_timeout() {
/*gateway_auth_manager*/ None,
);
let error = endpoint
.list_models(
tokio::time::pause();
let request = tokio::time::timeout(
MODELS_REFRESH_TIMEOUT * 2,
endpoint.list_models(
"0.0.0",
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
),
);
tokio::pin!(request);
let received = mock.wait_until_satisfied();
tokio::pin!(received);
let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 30);
// Keep the paused runtime runnable until the request reaches the server.
// Otherwise auto-advance can expire the catalog deadline during dispatch.
loop {
assert!(
std::time::Instant::now() < deadline,
"server did not receive the catalog request"
);
tokio::select! {
() = &mut received => break,
result = &mut request => {
panic!("catalog request finished before reaching the server: {result:?}");
}
() = tokio::task::yield_now() => {}
}
}
tokio::time::advance(MODELS_REFRESH_TIMEOUT + Duration::from_millis(/*millis*/ 1)).await;
let error = request
.await
.expect("catalog deadline should expire")
.expect_err("delayed catalog request should time out");
tokio::time::resume();
assert!(
matches!(error.details(), CodexErrorDetails::RequestTimeout),