test(kms): stabilize Vault transport retry coverage

Co-Authored-By: heihutu <heihutu@gmail.com>

crates/kms/src/backends/scripted_vault.rs

@@ -26,16 +26,16 @@ use std::sync::{Arc, Mutex};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
 
-/// One canned HTTP response.
-pub(crate) struct ScriptedResponse {
-    status: u16,
-    body: String,
+/// One scripted connection outcome.
+pub(crate) enum ScriptedResponse {
+    Http { status: u16, body: String },
+    Close,
 }
 
 impl ScriptedResponse {
     /// A 200 response carrying `data` inside the standard Vault envelope.
     pub(crate) fn ok(data: serde_json::Value) -> Self {
-        Self {
+        Self::Http {
             status: 200,
             body: serde_json::json!({
                 "request_id": "scripted",
@@ -50,11 +50,16 @@ impl ScriptedResponse {
 
     /// An error response in Vault's `{"errors": [...]}` format.
     pub(crate) fn error(status: u16, message: &str) -> Self {
-        Self {
+        Self::Http {
             status,
             body: serde_json::json!({ "errors": [message] }).to_string(),
         }
     }
+
+    /// Close the connection after consuming a request without sending an HTTP response.
+    pub(crate) fn close() -> Self {
+        Self::Close
+    }
 }
 
 /// A scripted stand-in Vault listening on a loopback port.
@@ -91,14 +96,14 @@ impl ScriptedVault {
                 let response = responses
                     .next()
                     .unwrap_or_else(|| ScriptedResponse::error(599, "scripted vault: script exhausted"));
-                let payload = format!(
-                    "HTTP/1.1 {} Scripted\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
-                    response.status,
-                    response.body.len(),
-                    response.body
-                );
-                let _ = stream.write_all(payload.as_bytes()).await;
-                let _ = stream.shutdown().await;
+                if let ScriptedResponse::Http { status, body } = response {
+                    let payload = format!(
+                        "HTTP/1.1 {status} Scripted\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
+                        body.len(),
+                    );
+                    let _ = stream.write_all(payload.as_bytes()).await;
+                    let _ = stream.shutdown().await;
+                }
             }
         });
 

crates/kms/src/backends/vault.rs

@@ -1336,6 +1336,8 @@ mod tests {
     use crate::backends::scripted_vault::{ScriptedResponse, ScriptedVault};
     use crate::config::{VaultAuthMethod, VaultConfig};
 
+    const SCRIPTED_RETRY_ATTEMPTS: u32 = 3;
+
     /// Vault + KMS config pair pointing at a scripted loopback Vault.
     fn scripted_configs(address: &str) -> (VaultConfig, KmsConfig) {
         let vault_config = VaultConfig {
@@ -1351,7 +1353,7 @@ mod tests {
         };
         let kms_config = KmsConfig {
             timeout: Duration::from_secs(5),
-            retry_attempts: 3,
+            retry_attempts: SCRIPTED_RETRY_ATTEMPTS,
             ..KmsConfig::default()
         };
         (vault_config, kms_config)
@@ -1420,6 +1422,31 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn wired_read_retries_closed_connections_within_budget() {
+        let (vault, client) = scripted_client(vec![ScriptedResponse::close(), ScriptedResponse::close()]).await;
+
+        let error = client
+            .get_key_data("wired-key")
+            .await
+            .expect_err("closed connections must exhaust the retry budget");
+        assert!(matches!(error, KmsError::BackendError { .. }), "got {error:?}");
+
+        let requests = vault.requests();
+        let expected_requests = usize::try_from(SCRIPTED_RETRY_ATTEMPTS).expect("retry attempts must fit usize");
+        assert_eq!(
+            requests.len(),
+            expected_requests,
+            "all budgeted retry attempts must reach Vault: {requests:?}"
+        );
+        assert!(
+            requests
+                .iter()
+                .all(|line| line == "GET /v1/secret/data/rustfs/kms/keys/wired-key"),
+            "all attempts must hit the same read endpoint: {requests:?}"
+        );
+    }
+
     #[tokio::test]
     async fn wired_read_does_not_retry_permission_errors() {
         let (vault, client) = scripted_client(vec![ScriptedResponse::error(403, "permission denied")]).await;

crates/kms/tests/vault_fault_injection.rs

@@ -14,8 +14,8 @@
 
 //! Fault-injection matrix for the Vault backend operation policy.
 //!
-//! Offline cases run against locally injected transport faults (a closed
-//! port, a listener that never responds) — deterministic, no external
+//! Offline cases run against locally injected transport faults (a listener
+//! that never responds) — deterministic, no external
 //! dependencies. Real-Vault cases are `#[ignore]`d and need a dev Vault
 //! (default `http://127.0.0.1:8200`, override with `RUSTFS_KMS_VAULT_ADDR`).
 //!
@@ -118,44 +118,6 @@ fn counter_value(snapshot: &[MetricEntry], name: &str, labels: &[(&str, &str)])
         .sum()
 }
 
-/// Connection refused: connection-class failures are retried up to the
-/// configured budget, then surface as a backend error.
-#[test]
-fn connection_refused_is_retried_within_budget() {
-    // Reserve a loopback port and release it so nothing is listening there.
-    let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("reserve a loopback port");
-    let address = format!("http://{}", listener.local_addr().expect("reserved port addr"));
-    drop(listener);
-
-    let snapshot = record_metrics(|| {
-        Box::pin(async move {
-            let client = VaultKmsBackend::new(kms_config(vault_config(&address, "unused"), Duration::from_secs(2), 2))
-                .await
-                .expect("client construction performs no network calls");
-            let error = KmsBackendTrait::describe_key(&client, describe_key_request("fault-injection-refused"))
-                .await
-                .expect_err("a refused connection must fail the operation");
-            assert!(matches!(error, KmsError::BackendError { .. }), "got {error:?}");
-        })
-    });
-
-    assert_eq!(
-        counter_value(&snapshot, ATTEMPT_FAILURES_TOTAL, &[("error_class", "retryable_conn")]),
-        2,
-        "both budgeted attempts must observe the refused connection"
-    );
-    assert_eq!(counter_value(&snapshot, OPERATIONS_TOTAL, &[("outcome", "budget_exhausted")]), 1);
-    // The static-token login records its own success; the Vault read must not.
-    assert_eq!(
-        counter_value(
-            &snapshot,
-            OPERATIONS_TOTAL,
-            &[("operation", "vault_kv2_read_key"), ("outcome", "success")]
-        ),
-        0
-    );
-}
-
 /// Stalled connection: a server that accepts but never responds is cut off by
 /// the per-attempt timeout (either the policy timer or the equally sized HTTP
 /// client timeout, whichever fires first) instead of hanging forever.