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_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,7 +4,7 @@
 #![doc = include_str!("README.md")]
 #![cfg(not(target_arch = "wasm32"))]
 
-use crate::TestContext;
+use crate::{Recording, TestContext};
 use azure_core::{
     error::{ErrorKind, ResultExt},
     time::Duration,
@@ -298,7 +298,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 +479,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 +498,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 +543,15 @@ impl PerfRunner {
     }
 }
 
+/// Instrument a client options appropriately for a perf test.
+///
+/// # Arguments:
+///
+/// - recording: Test recording.
+pub fn instrument_perf(recording: &Recording) -> azure_core::Result<()> {
+    recording.remove_recording(false)
+}
+
 #[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,22 @@ impl Policy for RecordingPolicy {
 #[derive(Debug, Default)]
 pub struct RecordingOptions {
     pub skip: Option<Skip>,
+    pub remove_recording: Option<bool>,
 }
 
 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 = if let Some(remove) = self.remove_recording {
+            RemoveRecording(remove).as_headers()
+        } else {
+            Ok(Vec::new().into_iter())
+        }?;
+
+        Ok(it_self.chain(it_remove).collect::<Vec<_>>().into_iter())
     }
 }

sdk/core/azure_core_test/src/recording.rs

@@ -116,7 +116,7 @@ impl Recording {
     ///     let recording = ctx.recording();
     ///
     ///     let mut options = MyClientOptions::default();
-    ///     ctx.instrument(&mut options.client_options);
+    ///     recording.instrument(&mut options.client_options);
     ///
     ///     let client = MyClient::new("https://azure.net", Some(options));
     ///     client.invoke().await
@@ -305,6 +305,11 @@ impl Recording {
         Ok(SkipGuard(self))
     }
 
+    pub(crate) fn remove_recording(&self, remove: bool) -> azure_core::Result<()> {
+        self.set_remove_recording(Some(remove))?;
+        Ok(())
+    }
+
     /// Gets the current [`TestMode`].
     pub fn test_mode(&self) -> TestMode {
         self.test_mode
@@ -461,6 +466,20 @@ impl Recording {
         Ok(())
     }
 
+    fn set_remove_recording(&self, remove: Option<bool>) -> azure_core::Result<()> {
+        let Some(policy) = self.recording_policy.get() else {
+            return Ok(());
+        };
+
+        let mut options = policy
+            .options
+            .write()
+            .map_err(|err| azure_core::Error::with_message(ErrorKind::Other, err.to_string()))?;
+        options.remove_recording = remove;
+
+        Ok(())
+    }
+
     /// Starts recording or playback.
     ///
     /// If playing back a recording, environment variable that were recorded will be reloaded.
@@ -587,6 +606,24 @@ 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 {
+        HeaderName::from_static("x-recording-remove")
+    }
+
+    fn value(&self) -> HeaderValue {
+        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/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_keys/perf/create_key.rs

@@ -4,7 +4,7 @@
 use crate::ENV_NAME;
 use azure_core::Result;
 use azure_core_test::{
-    perf::{CreatePerfTestReturn, PerfRunner, PerfTest, PerfTestMetadata, PerfTestOption},
+    perf::{self, CreatePerfTestReturn, PerfRunner, PerfTest, PerfTestMetadata, PerfTestOption},
     Recording, TestContext,
 };
 use azure_security_keyvault_keys::{
@@ -75,6 +75,7 @@ impl PerfTest for CreateKey {
 
         let mut client_options = KeyClientOptions::default();
         recording.instrument(&mut client_options.client_options);
+        perf::instrument_perf(recording)?;
 
         let vault_url = self
             .vault_url

sdk/keyvault/azure_security_keyvault_keys/perf/get_key.rs

@@ -4,7 +4,7 @@
 use crate::ENV_NAME;
 use azure_core::Result;
 use azure_core_test::{
-    perf::{CreatePerfTestReturn, PerfRunner, PerfTest, PerfTestMetadata, PerfTestOption},
+    perf::{self, CreatePerfTestReturn, PerfRunner, PerfTest, PerfTestMetadata, PerfTestOption},
     Recording, TestContext,
 };
 use azure_security_keyvault_keys::{
@@ -75,6 +75,7 @@ impl PerfTest for GetKey {
 
         let mut client_options = KeyClientOptions::default();
         recording.instrument(&mut client_options.client_options);
+        perf::instrument_perf(recording)?;
 
         let vault_url = self
             .vault_url

sdk/keyvault/azure_security_keyvault_secrets/perf/get_secret.rs

@@ -10,14 +10,14 @@
 //!
 //! To run the test, use the following command line arguments:
 //!
-//! cargo test --package azure_security_keyvault_secrets --test performance_tests -- --duration 10 --parallel 20 get_secret -u https://<my_vault>.vault.azure.net/
+//! cargo test --package azure_security_keyvault_secrets --test perf -- --duration 10 --parallel 20 get_secret -u https://<my_vault>.vault.azure.net/
 //!
 
 use std::sync::{Arc, OnceLock};
 
 use azure_core::Result;
 use azure_core_test::{
-    perf::{CreatePerfTestReturn, PerfRunner, PerfTest, PerfTestMetadata, PerfTestOption},
+    perf::{self, CreatePerfTestReturn, PerfRunner, PerfTest, PerfTestMetadata, PerfTestOption},
     Recording, TestContext,
 };
 use azure_security_keyvault_secrets::{
@@ -80,6 +80,7 @@ impl PerfTest for GetSecrets {
 
         let mut client_options = SecretClientOptions::default();
         recording.instrument(&mut client_options.client_options);
+        perf::instrument_perf(recording)?;
 
         let vault_url = self
             .vault_url

sdk/storage/azure_storage_blob/perf/list_blob_test.rs

@@ -110,7 +110,10 @@ impl PerfTest for ListBlobTest {
 
         let mut iterator = self.client.get().unwrap().list_blobs(None)?;
         while let Some(blob_segment) = iterator.try_next().await? {
-            let _body = blob_segment.into_body()?;
+            let body = blob_segment.into_body()?;
+            for blob in body.segment.blob_items.iter() {
+                std::hint::black_box(blob);
+            }
         }
 
         Ok(())