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:
@@ -43,64 +49,17 @@ A perf test has a name (`get_secret`, `list_blobs`, `upload_blob`, etc), a short
 
 Each perf test also has a set of command line options that are specific to the individual test, these are defined by a `PerfTestOptions` structure. It contains fields like help text for the option, activators
 
-Here is an example of test metadata for a performance test:
-
-```rust
-PerfTestMetadata {
-    name: "get_secret",
-    description: "Get a secret from Key Vault",
-    options: vec![PerfTestOption {
-        name: "vault_url",
-        display_message: "The URL of the Key Vault to use in the test",
-        mandatory: true,
-        short_activator: 'u',
-        long_activator: "vault-url",
-        expected_args_len: 1,
-        ..Default::default()
-    }],
-    create_test: Self::create_new_test,
-}
-```
+An example of perf test metadata [can be found here](https://github.com/Azure/azure-sdk-for-rust/blob/e47a38f93e7ac2797754c103da7fe8b177e46365/sdk/keyvault/azure_security_keyvault_keys/perf/create_key.rs#L26C1-L41C1)
 
-This defines a test named `get_secret` with a single required "vault_url" option.
-
-For this test, the `create_new_test` function looks like:
-
-```rust
-fn create_new_test(runner: PerfRunner) -> CreatePerfTestReturn {
-    async move {
-        let vault_url_ref: Option<&String> = runner.try_get_test_arg("vault_url")?;
-        let vault_url = vault_url_ref
-            .expect("vault_url argument is mandatory")
-            .clone();
-        Ok(Box::new(GetSecrets {
-            vault_url,
-            random_key_name: OnceLock::new(),
-            client: OnceLock::new(),
-        }) as Box<dyn PerfTest>)
-    }
-    .boxed()
-}
-```
+This defines a test named `create_key` with a single required "vault_url" option.
+
+An example of the `create_new_test` function [can be found here](https://github.com/Azure/azure-sdk-for-rust/blob/e47a38f93e7ac2797754c103da7fe8b177e46365/sdk/keyvault/azure_security_keyvault_keys/perf/get_key.rs#L42-L58)
 
 ### Test invocation
 
 The final piece of code which is necessary to run the performance tests is logic to hook up the tests with a test runner.
 
-```rust
-#[tokio::main]
-async fn main() -> azure_core::Result<()> {
-    let runner = PerfRunner::new(
-        env!("CARGO_MANIFEST_DIR"),
-        file!(),
-        vec![GetSecrets::test_metadata()],
-    )?;
-
-    runner.run().await?;
-
-    Ok(())
-}
-```
+An example of this, from the Keyvault Keys performance tests [can be found here](https://github.com/Azure/azure-sdk-for-rust/blob/e47a38f93e7ac2797754c103da7fe8b177e46365/sdk/keyvault/azure_security_keyvault_keys/perf/perf_tests.rs#L24-L35)
 
 This declares a perf test runner with a set of defined test metadata and runs the performance test. If your performance test suite has more than one performance test, then it should be added to the final parameter to the `PerfRunner::new()` function.
 

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

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#![doc = include_str!("README.md")]
 #![cfg(not(target_arch = "wasm32"))]
 
 use crate::TestContext;
@@ -298,7 +297,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 +478,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 +497,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),

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<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 mut headers: Vec<_> = self.skip.as_headers()?.collect();
+        if let Some(remove) = self.remove_recording {
+            headers.extend(RemoveRecording(remove).as_headers()?);
+        }
+        Ok(headers.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
@@ -157,6 +157,48 @@ impl Recording {
         options.per_try_policies.push(recording_policy);
     }
 
+    /// Update a recording with settings appropriate for a performance test.
+    ///
+    /// Instruments the [`ClientOptions`] to support recording and playing back of session records.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use azure_core_test::{recorded, perf::PerfTest, TestContext};
+    /// # use std::sync::{OnceLock, Arc};
+    /// # struct MyServiceClient;
+    /// # impl MyServiceClient {
+    /// #   fn new(endpoint: impl AsRef<str>, options: Option<MyServiceClientOptions>) -> Self { todo!() }
+    /// #   async fn invoke(&self) -> azure_core::Result<()> { todo!() }
+    /// # }
+    /// # #[derive(Default)]
+    /// # struct MyServiceClientOptions { client_options: azure_core::http::ClientOptions };
+    /// # #[derive(Default)]
+    /// # struct MyPerfTest { client: OnceLock<MyServiceClient> };
+    /// #[async_trait::async_trait]
+    /// impl PerfTest for MyPerfTest {
+    ///   async fn setup(&self, ctx: Arc<TestContext>) -> azure_core::Result<()> {
+    ///     let recording = ctx.recording();
+    ///
+    ///     let mut options = MyServiceClientOptions::default();
+    ///     recording.instrument_perf(&mut options.client_options)?;
+    ///
+    ///     let client = MyServiceClient::new("https://azure.net", Some(options));
+    ///     client.invoke().await
+    ///   }
+    ///   async fn run(&self, ctx: Arc<TestContext>) -> azure_core::Result<()>{ todo!()}
+    ///   async fn cleanup(&self, ctx: Arc<TestContext>) -> azure_core::Result<()>{ todo!()}
+    /// }
+    /// ```
+    ///
+    /// Note that this function is a no-op for live tests - it only affects recorded tests
+    /// in playback mode.
+    ///
+    pub fn instrument_perf(&self, options: &mut ClientOptions) -> azure_core::Result<()> {
+        self.instrument(options);
+        self.remove_recording(false)
+    }
+
     /// Get random data from the OS or recording.
     ///
     /// This will always be the OS cryptographically secure pseudo-random number generator (CSPRNG) when running live.
@@ -305,6 +347,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 +508,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 +648,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

@@ -74,7 +74,7 @@ impl PerfTest for CreateKey {
         let credential = recording.credential();
 
         let mut client_options = KeyClientOptions::default();
-        recording.instrument(&mut client_options.client_options);
+        recording.instrument_perf(&mut client_options.client_options)?;
 
         let vault_url = self
             .vault_url

sdk/keyvault/azure_security_keyvault_keys/perf/get_key.rs

@@ -74,7 +74,7 @@ impl PerfTest for GetKey {
         let credential = recording.credential();
 
         let mut client_options = KeyClientOptions::default();
-        recording.instrument(&mut client_options.client_options);
+        recording.instrument_perf(&mut client_options.client_options)?;
 
         let vault_url = self
             .vault_url

sdk/keyvault/azure_security_keyvault_secrets/perf/get_secret.rs

@@ -10,7 +10,7 @@
 //!
 //! 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};
@@ -79,7 +79,7 @@ impl PerfTest for GetSecrets {
         let credential = recording.credential();
 
         let mut client_options = SecretClientOptions::default();
-        recording.instrument(&mut client_options.client_options);
+        recording.instrument_perf(&mut client_options.client_options)?;
 
         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(())