Added test proxy support

sdk/core/azure_core/src/test.rs

@@ -31,6 +31,11 @@ impl TestMode {
     pub fn current() -> typespec::Result<Self> {
         std::env::var("AZURE_TEST_MODE").map_or_else(|_| Ok(TestMode::default()), |v| v.parse())
     }
+
+    /// Gets the `TestMode` from the `AZURE_TEST_MODE` environment variable or returns `None` if undefined.
+    pub fn current_opt() -> typespec::Result<Option<Self>> {
+        std::env::var("AZURE_TEST_MODE").map_or_else(|_| Ok(None), |v| v.parse().map(Some))
+    }
 }
 
 impl fmt::Debug for TestMode {

sdk/core/azure_core_opentelemetry/tests/telemetry_service_implementation.rs

@@ -257,7 +257,7 @@ async fn test_service_client_new(ctx: TestContext) -> Result<()> {
     let mut options = TestServiceClientOptions {
         ..Default::default()
     };
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = TestServiceClient::new(endpoint, credential, Some(options)).unwrap();
     assert_eq!(client.endpoint().as_str(), "https://www.microsoft.com/");
@@ -273,7 +273,7 @@ async fn test_service_client_get(ctx: TestContext) -> Result<()> {
     let endpoint = "https://azuresdkforcpp.azurewebsites.net";
     let credential = recording.credential().clone();
     let mut options = TestServiceClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = TestServiceClient::new(endpoint, credential, Some(options)).unwrap();
     let response = client.get("get", None).await;
@@ -305,7 +305,7 @@ async fn test_service_client_get_with_tracing(ctx: TestContext) -> Result<()> {
         },
         ..Default::default()
     };
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = TestServiceClient::new(endpoint, credential, Some(options)).unwrap();
     let response = client.get("get", None).await;
@@ -360,7 +360,7 @@ async fn test_service_client_get_tracing_error(ctx: TestContext) -> Result<()> {
         },
         ..Default::default()
     };
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = TestServiceClient::new(endpoint, credential.clone(), Some(options)).unwrap();
     let response = client.get("failing_url", None).await;
@@ -421,7 +421,7 @@ async fn test_service_client_get_with_function_tracing(ctx: TestContext) -> Resu
         },
         ..Default::default()
     };
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = TestServiceClient::new(endpoint, credential, Some(options)).unwrap();
     let response = client.get_with_function_tracing("get", None).await;
@@ -484,7 +484,7 @@ async fn test_service_client_get_with_function_tracing_error(ctx: TestContext) -
         },
         ..Default::default()
     };
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = TestServiceClient::new(endpoint, credential, Some(options)).unwrap();
     let response = client.get_with_function_tracing("failing_url", None).await;

sdk/core/azure_core_opentelemetry/tests/telemetry_service_macros.rs

@@ -220,7 +220,7 @@ mod tests {
             },
             ..Default::default()
         };
-        recording.instrument(&mut options.client_options);
+        recording.instrument(&mut options.client_options, None);
 
         TestServiceClientWithMacros::new(endpoint, credential, Some(options)).unwrap()
     }
@@ -284,7 +284,7 @@ mod tests {
         let mut options = TestServiceClientWithMacrosOptions {
             ..Default::default()
         };
-        recording.instrument(&mut options.client_options);
+        recording.instrument(&mut options.client_options, None);
 
         let client = TestServiceClientWithMacros::new(endpoint, credential, Some(options)).unwrap();
         assert_eq!(client.endpoint().as_str(), "https://microsoft.com/");

sdk/core/azure_core_test/src/lib.rs

@@ -17,6 +17,7 @@ pub mod stream;
 pub mod tracing;
 use azure_core::Error;
 pub use azure_core::{error::ErrorKind, test::TestMode};
+pub use proxy::policy::RecordingOptions;
 pub use proxy::{matchers::*, sanitizers::*};
 pub use recording::*;
 use std::path::{Path, PathBuf};

sdk/core/azure_core_test/src/perf/README.md

@@ -25,6 +25,12 @@ Each performance test has the following standardized parameters:
 
 Each test has its own set of parameters which are specific to the test.
 
+***NOTE: Performance Tests are "recorded" tests***
+
+This means that they follow the same rules as tests annotated with the `#[recorded::test]` attribute. There is one difference between perf tests and tests with the `recorded::test` attribute: perf tests default to `live` mode, and normal `recorded::test` tests default to `playback` mode.
+
+To configure the tests for record mode tests, set `AZURE_TEST_MODE` to `record` before running your performance tests, and to run your tests using the test proxy, set `AZURE_TEST_MODE` to `playback`
+
 ## Test authoring
 
 Performance tests have three phases:

sdk/core/azure_core_test/src/perf/mod.rs

@@ -4,9 +4,10 @@
 #![doc = include_str!("README.md")]
 #![cfg(not(target_arch = "wasm32"))]
 
-use crate::TestContext;
+use crate::{Recording, RecordingOptions, RemoveRecording, TestContext};
 use azure_core::{
     error::{ErrorKind, ResultExt},
+    http::ClientOptions,
     time::Duration,
     Error, Result,
 };
@@ -298,7 +299,7 @@ impl PerfRunner {
         let test_instance = (test.create_test)(self.clone()).await?;
         let test_instance: Arc<dyn PerfTest> = Arc::from(test_instance);
 
-        let test_mode = crate::TestMode::current()?;
+        let test_mode = crate::TestMode::current_opt()?.unwrap_or(crate::TestMode::Live);
 
         let context = Arc::new(
             crate::recorded::start(
@@ -479,10 +480,17 @@ impl PerfRunner {
                     .value_parser(clap::value_parser!(u32))
                     .global(false),
             )
-            .arg(clap::arg!(--sync).global(true).required(false))
+            .arg(clap::arg!(--sync "Run synchronous tests (ignored)")
+                .global(true)
+                .required(false))
+            .arg(clap::arg!(--"test-proxy" <URL> "The URL of the test proxy, ignored.")
+                .global(true)
+                .value_parser(clap::value_parser!(String))
+                .required(false))
             .arg(
                 clap::arg!(--parallel <COUNT> "The number of concurrent tasks to use when running each test")
                     .required(false)
+                    .short('p')
                     .default_value("1")
                     .value_parser(clap::value_parser!(u32))
                     .global(true),
@@ -491,6 +499,7 @@ impl PerfRunner {
             .arg(
                 clap::arg!(--duration <SECONDS> "The duration of each test in seconds")
                     .required(false)
+                    .short('d')
                     .default_value("30")
                     .value_parser(clap::value_parser!(i64))
                     .global(true),
@@ -535,6 +544,23 @@ impl PerfRunner {
     }
 }
 
+/// Instrument a client options appropriately for a perf test.
+///
+/// # Arguments:
+///
+/// - recording: Test recording.
+/// - client_options: Core client options
+///
+pub fn instrument_for_perf_test(recording: &Recording, client_options: &mut ClientOptions) {
+    recording.instrument(
+        client_options,
+        Some(RecordingOptions {
+            remove_recording: Some(RemoveRecording(false)),
+            ..Default::default()
+        }),
+    );
+}
+
 #[cfg(test)]
 mod config_tests;
 

sdk/core/azure_core_test/src/proxy/policy.rs

@@ -3,7 +3,7 @@
 
 use crate::{
     proxy::{RecordingId, RECORDING_MODE, RECORDING_UPSTREAM_BASE_URI},
-    Skip,
+    RemoveRecording, Skip,
 };
 use async_trait::async_trait;
 use azure_core::{
@@ -118,13 +118,18 @@ impl Policy for RecordingPolicy {
 #[derive(Debug, Default)]
 pub struct RecordingOptions {
     pub skip: Option<Skip>,
+    pub remove_recording: Option<RemoveRecording>,
 }
 
 impl AsHeaders for RecordingOptions {
     type Error = Infallible;
     type Iter = std::vec::IntoIter<(HeaderName, HeaderValue)>;
 
     fn as_headers(&self) -> Result<Self::Iter, Self::Error> {
-        self.skip.as_headers()
+        let it_self = self.skip.as_headers()?;
+
+        let it_remove = self.remove_recording.as_headers()?;
+
+        Ok(it_self.chain(it_remove).collect::<Vec<_>>().into_iter())
     }
 }

sdk/core/azure_core_test/src/recording.rs

@@ -13,7 +13,7 @@ use crate::{
             ClientAddSanitizerOptions, ClientRemoveSanitizersOptions, ClientSetMatcherOptions,
         },
         models::{SanitizerList, StartPayload, VariablePayload},
-        policy::RecordingPolicy,
+        policy::{RecordingOptions, RecordingPolicy},
         Proxy, ProxyExt, RecordingId,
     },
     recording::policy::RecordingModePolicy,
@@ -40,7 +40,7 @@ use std::{
     env,
     sync::{Arc, Mutex, OnceLock, RwLock},
 };
-use tracing::span::EnteredSpan;
+use tracing::{span::EnteredSpan, trace};
 
 /// Represents a playback or recording session using the [`Proxy`].
 #[derive(Debug)]
@@ -116,13 +116,17 @@ impl Recording {
     ///     let recording = ctx.recording();
     ///
     ///     let mut options = MyClientOptions::default();
-    ///     ctx.instrument(&mut options.client_options);
+    ///     ctx.instrument(&mut options.client_options, None);
     ///
     ///     let client = MyClient::new("https://azure.net", Some(options));
     ///     client.invoke().await
     /// }
     /// ```
-    pub fn instrument(&self, options: &mut ClientOptions) {
+    pub fn instrument(
+        &self,
+        options: &mut ClientOptions,
+        recording_options: Option<RecordingOptions>,
+    ) {
         let Some(client) = self.proxy.client() else {
             return;
         };
@@ -142,14 +146,15 @@ impl Recording {
             options.per_call_policies.push(test_mode_policy);
         }
 
+        trace!("Setting recording options to {recording_options:?}");
         let recording_policy = self
             .recording_policy
             .get_or_init(|| {
                 Arc::new(RecordingPolicy {
                     test_mode: self.test_mode,
                     host: Some(client.endpoint().clone()),
                     recording_id: self.id.clone(),
-                    ..Default::default()
+                    options: RwLock::new(recording_options.unwrap_or_default()),
                 })
             })
             .clone();
@@ -587,6 +592,26 @@ impl Drop for SkipGuard<'_> {
     }
 }
 
+/// Whether to remove records during recording playback.
+///
+/// This option is used for test recordings, if true, the recording will be removed from the test-proxy when retrieved,
+/// otherwise it will be kept. The default is true.
+///
+#[derive(Debug)]
+pub struct RemoveRecording(pub bool);
+
+impl Header for RemoveRecording {
+    fn name(&self) -> HeaderName {
+        trace!("RemoveRecording::name");
+        HeaderName::from_static("x-recording-remove")
+    }
+
+    fn value(&self) -> HeaderValue {
+        trace!("RemoveRecording::value: {}", self.0);
+        HeaderValue::from_static(if self.0 { "true" } else { "false" })
+    }
+}
+
 /// Options for getting variables from a [`Recording`].
 #[derive(Clone, Debug)]
 pub struct VarOptions {

sdk/core/typespec_client_core/src/http/pipeline.rs

@@ -76,11 +76,11 @@ impl Pipeline {
         let retry_policy = options.retry.to_policy(pipeline_options.retry_headers);
         pipeline.push(retry_policy);
 
-        pipeline.push(Arc::new(LoggingPolicy::new(options.logging)));
-
         pipeline.extend_from_slice(&per_try_policies);
         pipeline.extend_from_slice(&options.per_try_policies);
 
+        pipeline.push(Arc::new(LoggingPolicy::new(options.logging)));
+
         let transport: Arc<dyn Policy> =
             Arc::new(TransportPolicy::new(options.transport.unwrap_or_default()));
         pipeline.push(transport);

sdk/cosmos/azure_data_cosmos/tests/framework/test_account.rs

@@ -111,7 +111,7 @@ impl TestAccount {
 
         self.context
             .recording()
-            .instrument(&mut options.client_options);
+            .instrument(&mut options.client_options, None);
 
         Ok(azure_data_cosmos::CosmosClient::with_key(
             &self.endpoint,

sdk/eventhubs/azure_messaging_eventhubs_checkpointstore_blob/tests/checkpoint_unit_tests.rs

@@ -14,7 +14,7 @@ use tracing::trace;
 pub fn create_test_checkpoint_store(recording: &Recording) -> Result<Arc<BlobCheckpointStore>> {
     let credential = recording.credential();
     let mut options = BlobContainerClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
     let blob_container_client = BlobContainerClient::new(
         &recording.var("AZURE_STORAGE_BLOB_ENDPOINT", None),
         recording.var("AZURE_STORAGE_BLOB_CONTAINER", None),

sdk/keyvault/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "rust",
   "TagPrefix": "rust/keyvault",
-  "Tag": "rust/keyvault_3dde96c6e5"
+  "Tag": "rust/keyvault_568367d8ae"
 }

sdk/keyvault/azure_security_keyvault_certificates/tests/certificate_client.rs

@@ -41,7 +41,7 @@ async fn certificate_roundtrip(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),
@@ -78,7 +78,7 @@ async fn update_certificate_properties(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),
@@ -138,7 +138,7 @@ async fn list_certificates(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),
@@ -181,7 +181,7 @@ async fn purge_certificate(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),
@@ -237,7 +237,7 @@ async fn sign_jwt_with_ec_certificate(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),
@@ -275,7 +275,7 @@ async fn sign_jwt_with_ec_certificate(ctx: TestContext) -> Result<()> {
         .await?;
 
     let mut key_options = KeyClientOptions::default();
-    recording.instrument(&mut key_options.client_options);
+    recording.instrument(&mut key_options.client_options, None);
 
     // Sign a JWT.
     let key_client = KeyClient::new(
@@ -310,7 +310,7 @@ async fn get_certificate_operation(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),
@@ -358,7 +358,7 @@ async fn create_invalid_certificate(ctx: TestContext) -> Result<()> {
     recording.remove_sanitizers(&[SANITIZE_BODY_NAME]).await?;
 
     let mut options = CertificateClientOptions::default();
-    recording.instrument(&mut options.client_options);
+    recording.instrument(&mut options.client_options, None);
 
     let client = CertificateClient::new(
         recording.var("AZURE_KEYVAULT_URL", None).as_str(),