Add ability to specify path as regex in snapshots

Cargo.lock

@@ -377,6 +377,7 @@ dependencies = [
  "serde_derive_default",
  "serde_json",
  "serde_json_bytes",
+ "serde_regex",
  "serde_urlencoded",
  "serde_yaml",
  "serial_test",
@@ -6284,6 +6285,16 @@ dependencies = [
  "thiserror 1.0.69",
 ]
 
+[[package]]
+name = "serde_regex"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a8136f1a4ea815d7eac4101cfd0b16dc0cb5e1fe1b8609dfd728058656b7badf"
+dependencies = [
+ "regex",
+ "serde",
+]
+
 [[package]]
 name = "serde_urlencoded"
 version = "0.7.1"

apollo-router/Cargo.toml

@@ -53,7 +53,7 @@ telemetry_next = []
 ci = []
 
 # Enables the HTTP snapshot server for testing
-snapshot = ["axum-server"]
+snapshot = ["axum-server", "serde_regex"]
 
 [package.metadata.docs.rs]
 features = ["docs_rs"]
@@ -203,6 +203,7 @@ serde.workspace = true
 serde_derive_default = "0.1"
 serde_json_bytes.workspace = true
 serde_json.workspace = true
+serde_regex = { version = "1.1.0", optional = true }
 serde_urlencoded = "0.7.1"
 serde_yaml = "0.8.26"
 static_assertions = "1.1.0"
@@ -313,6 +314,7 @@ rhai = { version = "1.17.1", features = [
     "internals",
     "testing-environ",
 ] }
+serde_regex = { version = "1.1.0" }
 serial_test = { version = "3.1.1" }
 tempfile.workspace = true
 test-log = { version = "0.2.16", default-features = false, features = [

apollo-router/src/test_harness/http_snapshot.rs

@@ -18,12 +18,23 @@
 //!
 //! Any requests made to the snapshot server will be proxied on to the given base URL, and the
 //! responses will be saved to the given file. The next time the snapshot server receives the
-//! same request (same relative path, HTTP method, and request body), it will respond with the
-//! response recorded in the file rather than sending the request to the upstream server.
+//! same request, it will respond with the response recorded in the file rather than sending the
+//! request to the upstream server.
 //!
 //! The snapshot file can be manually edited to manipulate responses for testing purposes, or to
 //! redact information that you don't want to include in source-controlled snapshot files.
 //!
+//! Requests are matched to snapshots based on the URL path, HTTP method, and base64 encoding of
+//! the request body (if there is one). If the snapshot specifies the `path` field, the URL path
+//! must match exactly. Alternatively, a snapshot containing a `regex` field will match any URL
+//! path that matches the regular expression. A snapshot with an exact `path` match takes
+//! precedence over a snapshot with `regex`. Snapshots recorded by the server will always specify
+//! `path`. The only way to use `regex` is to manually edit a snapshot file. A typical pattern is
+//! to record a snapshot from a REST API, then manually change `path` to `regex` and replace the
+//! variable part of the path with `.*`. Note that any special characters in the path that have
+//! meaning to the `regex` crate must be escaped with `\\`, such as the `?` in a URL query
+//! parameter.
+//!
 //! The offline mode will never call the upstream server, and will always return a saved snapshot
 //! response. If one is not available, a `500` error is returned. This is useful for tests, for
 //! example to ensure that CI builds never attempt to access the network.
@@ -36,10 +47,7 @@
 //! This is typically desirable, as headers may contain ephemeral information like dates or tokens.
 //!
 //! **IMPORTANT:** this module stores HTTP responses to the local file system in plain text. It
-//! should not be used with production APIs that return sensitive data.
-//!
-//! This module should also not be used in conjunction with performance testing, as returning
-//! snapshot data locally will be much faster than sending HTTP requests to an external server.
+//! should not be used with APIs that return sensitive data.
 
 use std::collections::BTreeMap;
 use std::net::SocketAddr;
@@ -69,6 +77,7 @@ use hyper::StatusCode;
 use hyper_rustls::ConfigBuilderExt;
 use indexmap::IndexMap;
 use parking_lot::Mutex;
+use regex::Regex;
 use serde::Deserialize;
 use serde::Serialize;
 use serde_json_bytes::json;
@@ -99,11 +108,14 @@ static FILTERED_HEADERS: [HeaderName; 6] = [
 #[derive(Debug, thiserror::Error)]
 enum SnapshotError {
     /// Unable to load snapshots
-    #[error("unable to load snapshots")]
+    #[error("unable to load snapshot file - {0}")]
     IoError(#[from] std::io::Error),
     /// Unable to parse snapshots
-    #[error("unable to parse snapshots")]
+    #[error("unable to parse snapshots - {0}")]
     ParseError(#[from] serde_json::Error),
+    /// Invalid snapshot
+    #[error("invalid snapshot - {0}")]
+    InvalidSnapshot(String),
 }
 
 /// A server that mocks an API using snapshots recorded from actual HTTP responses.
@@ -118,7 +130,8 @@ pub struct SnapshotServer {
 struct SnapshotServerState {
     client: HttpClientService,
     base_url: Uri,
-    snapshots: Arc<Mutex<BTreeMap<String, Snapshot>>>,
+    snapshots_by_key: Arc<Mutex<BTreeMap<String, Snapshot>>>,
+    snapshots_by_regex: Arc<Mutex<Vec<Snapshot>>>,
     snapshot_file: Box<Path>,
     offline: bool,
     update: bool,
@@ -217,6 +230,7 @@ async fn handle(
                     request: Request {
                         method: Some(method.to_string()),
                         path: Some(path),
+                        regex: None,
                         body: request_json_body,
                     },
                     response: Response {
@@ -232,9 +246,14 @@ async fn handle(
                     },
                 };
                 {
-                    let mut snapshots = state.snapshots.lock();
-                    snapshots.insert(key, snapshot.clone());
-                    if let Err(e) = save(state.snapshot_file, &mut snapshots) {
+                    let mut snapshots_by_key = state.snapshots_by_key.lock();
+                    let mut snapshots_by_regex = state.snapshots_by_regex.lock();
+                    snapshots_by_key.insert(key, snapshot.clone());
+                    if let Err(e) = save(
+                        state.snapshot_file,
+                        &mut snapshots_by_key,
+                        &mut snapshots_by_regex,
+                    ) {
                         error!(
                             url = %uri,
                             method = %method,
@@ -263,23 +282,46 @@ fn response_from_snapshot(
     method: &Method,
     key: &String,
 ) -> Option<http::Response<RouterBody>> {
-    let mut snapshots = state.snapshots.lock();
+    let mut snapshots_by_key = state.snapshots_by_key.lock();
+    let snapshots_by_regex = state.snapshots_by_regex.lock();
     if state.update {
-        snapshots.remove(key);
+        snapshots_by_key.remove(key);
         None
     } else {
-        snapshots.get(key).and_then(|snapshot| {
-            debug!(
-                url = %uri,
-                method = %method,
-                "Found existing snapshot"
-            );
-            snapshot
-                .clone()
-                .into_response()
-                .map_err(|e| error!("Unable to convert snapshot into HTTP response: {:?}", e))
-                .ok()
-        })
+        snapshots_by_key
+            .get(key)
+            .inspect(|snapshot| {
+                debug!(
+                    url = %uri,
+                    method = %method,
+                    path = %snapshot.request.path.as_ref().unwrap_or(&String::from("")),
+                    "Found existing snapshot"
+                );
+            })
+            .or_else(|| {
+                // Look up snapshot using regex
+                for snapshot in snapshots_by_regex.iter() {
+                    if let Some(regex) = &snapshot.request.regex {
+                        if regex.is_match(uri) {
+                            debug!(
+                                url = %uri,
+                                method = %method,
+                                regex = %regex.to_string(),
+                                "Found existing snapshot"
+                            );
+                            return Some(snapshot);
+                        }
+                    }
+                }
+                None
+            })
+            .and_then(|snapshot| {
+                snapshot
+                    .clone()
+                    .into_response()
+                    .map_err(|e| error!("Unable to convert snapshot into HTTP response: {:?}", e))
+                    .ok()
+            })
     }
 }
 
@@ -320,24 +362,41 @@ fn map_headers<F: Fn(&str) -> bool>(
 
 fn save<P: AsRef<Path>>(
     path: P,
-    snapshots: &mut BTreeMap<String, Snapshot>,
+    snapshots_by_key: &mut BTreeMap<String, Snapshot>,
+    snapshots_by_regex: &mut [Snapshot],
 ) -> Result<(), SnapshotError> {
     let path = path.as_ref();
     if let Some(parent) = path.parent() {
         std::fs::create_dir_all(parent)?;
     }
-    let snapshots = snapshots.values().cloned().collect::<Vec<_>>();
+    let snapshots = snapshots_by_key
+        .values()
+        .cloned()
+        .chain(snapshots_by_regex.iter().cloned())
+        .collect::<Vec<_>>();
     std::fs::write(path, serde_json::to_string_pretty(&snapshots)?).map_err(Into::into)
 }
 
-fn load<P: AsRef<Path>>(path: P) -> Result<BTreeMap<String, Snapshot>, SnapshotError> {
+fn load<P: AsRef<Path>>(
+    path: P,
+) -> Result<(BTreeMap<String, Snapshot>, Vec<Snapshot>), SnapshotError> {
     let str = std::fs::read_to_string(path)?;
     let snapshots: Vec<Snapshot> = serde_json::from_str(&str)?;
-    info!("Loaded {} snapshots", snapshots.len());
-    Ok(snapshots
-        .into_iter()
-        .map(|snapshot| (snapshot.key(), snapshot))
-        .collect())
+    let mut snapshots_by_key: BTreeMap<String, Snapshot> = Default::default();
+    let mut snapshots_by_regex: Vec<Snapshot> = Default::default();
+    for snapshot in snapshots.into_iter() {
+        if snapshot.request.regex.is_some() {
+            if snapshot.request.path.is_some() {
+                return Err(SnapshotError::InvalidSnapshot(String::from(
+                    "snapshot cannot specify both regex and path",
+                )));
+            }
+            snapshots_by_regex.push(snapshot);
+        } else {
+            snapshots_by_key.insert(snapshot.key(), snapshot);
+        }
+    }
+    Ok((snapshots_by_key, snapshots_by_regex))
 }
 
 impl SnapshotServer {
@@ -412,12 +471,31 @@ impl SnapshotServer {
 
         let snapshot_file = snapshot_path.as_ref();
 
-        let snapshots = load(snapshot_file).unwrap_or_else(|_| {
-            if offline {
-                warn!("Unable to load snapshot file in offline mode - all requests will fail");
+        let (snapshots_by_key, snapshots_by_regex) = match load(snapshot_file) {
+            Err(SnapshotError::IoError(ioe)) if ioe.kind() == std::io::ErrorKind::NotFound => {
+                if offline {
+                    warn!("Snapshot file not found in offline mode - all requests will fail");
+                } else {
+                    info!("Snapshot file not found - new snapshot file will be recorded");
+                }
+                (BTreeMap::default(), vec![])
             }
-            BTreeMap::default()
-        });
+            Err(e) => {
+                if offline {
+                    warn!("Unable to load snapshot file in offline mode - all requests will fail: {e}");
+                } else {
+                    warn!("Unable to load snapshot file - new snapshot file will be recorded: {e}");
+                }
+                (BTreeMap::default(), vec![])
+            }
+            Ok((snapshots_by_key, snapshots_by_regex)) => {
+                info!(
+                    "Loaded {} snapshots",
+                    snapshots_by_key.len() + snapshots_by_regex.len()
+                );
+                (snapshots_by_key, snapshots_by_regex)
+            }
+        };
 
         let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
 
@@ -436,7 +514,8 @@ impl SnapshotServer {
             .with_state(SnapshotServerState {
                 client: http_service,
                 base_url: base_url.clone(),
-                snapshots: Arc::new(Mutex::new(snapshots.clone())),
+                snapshots_by_key: Arc::new(Mutex::new(snapshots_by_key)),
+                snapshots_by_regex: Arc::new(Mutex::new(snapshots_by_regex)),
                 snapshot_file: Box::from(snapshot_file),
                 offline,
                 update,
@@ -533,6 +612,8 @@ fn snapshot_key(method: Option<&str>, path: Option<&str>, body: &Value) -> Strin
 struct Request {
     method: Option<String>,
     path: Option<String>,
+    #[serde(with = "serde_regex", skip_serializing_if = "Option::is_none", default)]
+    regex: Option<Regex>,
     body: Value,
 }