add RuleBuilder::then_compute_output (#4299)

this allows the developer to compute a mocked output using content from
the input.

## Motivation and Context
Some richer use cases for stubbing/mocking require computing a stubbed
output using data from the input, for instance stubbing a method that
returns the item that was put.

## Description
Add a new method to `RuleBuilder` that allows for providing a function
that computes the output. That function takes in a reference to the
input as an argument. I'm not sure how this works with the mutable
inputs like streamed bytes, but seems to work for the simple stuff.

## Testing
added a simple unit test

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Co-authored-by: Aaron Todd <aajtodd@users.noreply.github.com>

Author: greenwoodcm

.changelog/1757712519.md

@@ -0,0 +1,19 @@
+---
+applies_to: ["client"]
+authors: ["greenwoodcm"]
+references: ["smithy-rs#4299"]
+breaking: false
+new_feature: true
+bug_fix: false
+---
+Added a new `then_compute_output` to `aws-smithy-mocks` rule builder that allows using the input type when computing a mocked response, e.g. 
+```rs
+// Return a computed output based on the input
+let compute_rule = mock!(Client::get_object)
+    .then_compute_output(|req| {
+        let key = req.key().unwrap_or("unknown");
+        GetObjectOutput::builder()
+            .body(ByteStream::from_static(format!("content for {}", key).as_bytes()))
+            .build()
+    });
+```

aws/sdk/integration-tests/s3/tests/mocks.rs

@@ -89,6 +89,39 @@ async fn test_mock_client() {
     assert_eq!(err.code(), Some("InvalidAccessKey"));
 }
 
+#[tokio::test]
+async fn test_mock_client_compute() {
+    let s3_computed = mock!(aws_sdk_s3::Client::get_object)
+        .match_requests(|inp| {
+            inp.bucket() == Some("test-bucket") && inp.key() == Some("correct-key")
+        })
+        .then_compute_output(|input| {
+            let content =
+                format!("{}.{}", input.bucket().unwrap(), input.key().unwrap()).into_bytes();
+            GetObjectOutput::builder()
+                .body(ByteStream::from(content))
+                .build()
+        });
+
+    let s3 = mock_client!(aws_sdk_s3, &[&s3_computed]);
+
+    let data = s3
+        .get_object()
+        .bucket("test-bucket")
+        .key("correct-key")
+        .send()
+        .await
+        .expect("success response")
+        .body
+        .collect()
+        .await
+        .expect("successful read")
+        .to_vec();
+
+    assert_eq!(data, b"test-bucket.correct-key");
+    assert_eq!(s3_computed.num_calls(), 1);
+}
+
 #[tokio::test]
 async fn test_mock_client_sequence() {
     let rule = mock!(aws_sdk_s3::Client::get_object)

rust-runtime/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "aws-smithy-mocks"
-version = "0.1.1"
+version = "0.1.2"
 authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
 description = "Testing utilities for smithy-rs generated clients"
 edition = "2021"
@@ -9,7 +9,7 @@ repository = "https://github.com/smithy-lang/smithy-rs"
 
 [dependencies]
 aws-smithy-types = { path = "../aws-smithy-types" }
-aws-smithy-runtime-api = { path = "../aws-smithy-runtime-api", features = ["client", "http-1x"] }
+aws-smithy-runtime-api = { path = "../aws-smithy-runtime-api", features = ["client", "http-1x", "test-util"] }
 aws-smithy-http-client = { path = "../aws-smithy-http-client", features = ["test-util"] }
 http = "1"
 

rust-runtime/aws-smithy-mocks/Cargo.toml

@@ -83,6 +83,15 @@ let http_rule = mock!(Client::get_object)
         StatusCode::try_from(503).unwrap(),
         SdkBody::from("service unavailable")
     ));
+
+// Return a computed output based on the input
+let compute_rule = mock!(Client::get_object)
+    .then_compute_output(|req| {
+        let key = req.key().unwrap_or("unknown");
+        GetObjectOutput::builder()
+            .body(ByteStream::from_static(format!("content for {}", key).as_bytes()))
+            .build()
+    });
 ```
 
 ### Response Sequences

rust-runtime/aws-smithy-mocks/README.md

@@ -9,7 +9,7 @@ use aws_smithy_runtime_api::box_error::BoxError;
 use aws_smithy_runtime_api::client::http::SharedHttpClient;
 use aws_smithy_runtime_api::client::interceptors::context::{
     BeforeSerializationInterceptorContextMut, BeforeTransmitInterceptorContextMut, Error,
-    FinalizerInterceptorContextMut, Output,
+    FinalizerInterceptorContextMut, Input, Output,
 };
 use aws_smithy_runtime_api::client::interceptors::Intercept;
 use aws_smithy_runtime_api::client::orchestrator::{HttpResponse, OrchestratorError};
@@ -132,7 +132,7 @@ impl Intercept for MockResponseInterceptor {
                     }
 
                     // Rule matches and is not exhausted, get the response
-                    if let Some(response) = rule.next_response() {
+                    if let Some(response) = rule.next_response(input) {
                         matching_rule = Some(rule.clone());
                         matching_response = Some(response);
                     } else {
@@ -154,7 +154,7 @@ impl Intercept for MockResponseInterceptor {
                     }
 
                     if (rule.matcher)(input) {
-                        if let Some(response) = rule.next_response() {
+                        if let Some(response) = rule.next_response(input) {
                             matching_rule = Some(rule.clone());
                             matching_response = Some(response);
                             break;
@@ -198,7 +198,10 @@ impl Intercept for MockResponseInterceptor {
         if active_response.is_none() {
             // in the case of retries we try to get the next response if it has been consumed
             if let Some(active_rule) = cfg.load::<ActiveRule>() {
-                let next_resp = active_rule.0.next_response();
+                // During retries, input is not available in modify_before_transmit.
+                // For HTTP status responses that don't use the input, we can use a dummy input.
+                let dummy_input = Input::doesnt_matter();
+                let next_resp = active_rule.0.next_response(&dummy_input);
                 active_response = next_resp;
             }
         }
@@ -444,6 +447,43 @@ mod tests {
         assert_eq!(rule.num_calls(), 3);
     }
 
+    #[tokio::test]
+    async fn test_compute_output() {
+        // Create a rule that computes its responses based off of input data
+        let rule = create_rule_builder()
+            .match_requests(|input| input.bucket == "test-bucket" && input.key == "test-key")
+            .then_compute_output(|input| TestOutput {
+                content: format!("{}.{}", input.bucket, input.key),
+            });
+
+        // Create an interceptor with the rule
+        let interceptor = MockResponseInterceptor::new()
+            .rule_mode(RuleMode::Sequential)
+            .with_rule(&rule);
+
+        let operation = create_test_operation(interceptor, true);
+
+        let result = operation
+            .invoke(TestInput::new("test-bucket", "test-key"))
+            .await;
+
+        // Should succeed with the output derived from input
+        assert!(
+            result.is_ok(),
+            "Expected success but got error: {:?}",
+            result.err()
+        );
+        assert_eq!(
+            result.unwrap(),
+            TestOutput {
+                content: "test-bucket.test-key".to_string()
+            }
+        );
+
+        // Verify the rule was used once, no retries
+        assert_eq!(rule.num_calls(), 1);
+    }
+
     #[should_panic(
         expected = "must_match was enabled but no rules matched or all rules were exhausted for"
     )]

rust-runtime/aws-smithy-mocks/src/interceptor.rs

@@ -32,7 +32,7 @@ pub(crate) enum MockResponse<O, E> {
 
 /// A function that matches requests.
 type MatchFn = Arc<dyn Fn(&Input) -> bool + Send + Sync>;
-type ServeFn = Arc<dyn Fn(usize) -> Option<MockResponse<Output, Error>> + Send + Sync>;
+type ServeFn = Arc<dyn Fn(usize, &Input) -> Option<MockResponse<Output, Error>> + Send + Sync>;
 
 /// A rule for matching requests and providing mock responses.
 ///
@@ -67,9 +67,10 @@ impl fmt::Debug for Rule {
 
 impl Rule {
     /// Creates a new rule with the given matcher, response handler, and max responses.
+    #[allow(clippy::type_complexity)]
     pub(crate) fn new<O, E>(
         matcher: MatchFn,
-        response_handler: Arc<dyn Fn(usize) -> Option<MockResponse<O, E>> + Send + Sync>,
+        response_handler: Arc<dyn Fn(usize, &Input) -> Option<MockResponse<O, E>> + Send + Sync>,
         max_responses: usize,
         is_simple: bool,
     ) -> Self
@@ -79,9 +80,9 @@ impl Rule {
     {
         Rule {
             matcher,
-            response_handler: Arc::new(move |idx: usize| {
+            response_handler: Arc::new(move |idx: usize, input: &Input| {
                 if idx < max_responses {
-                    response_handler(idx).map(|resp| match resp {
+                    response_handler(idx, input).map(|resp| match resp {
                         MockResponse::Output(o) => MockResponse::Output(Output::erase(o)),
                         MockResponse::Error(e) => MockResponse::Error(Error::erase(e)),
                         MockResponse::Http(http_resp) => MockResponse::Http(http_resp),
@@ -102,9 +103,9 @@ impl Rule {
     }
 
     /// Gets the next response.
-    pub(crate) fn next_response(&self) -> Option<MockResponse<Output, Error>> {
+    pub(crate) fn next_response(&self, input: &Input) -> Option<MockResponse<Output, Error>> {
         let idx = self.call_count.fetch_add(1, Ordering::SeqCst);
-        (self.response_handler)(idx)
+        (self.response_handler)(idx, input)
     }
 
     /// Returns the number of times this rule has been called.
@@ -227,9 +228,31 @@ where
     {
         self.sequence().http_response(response_fn).build_simple()
     }
+
+    /// Creates a rule that computes an output based on the input.
+    ///
+    /// This allows generating responses based on the input request.
+    ///
+    /// # Examples
+    ///
+    /// ```rust,ignore
+    /// let rule = mock!(Client::get_object)
+    ///     .compute_output(|req| {
+    ///         GetObjectOutput::builder()
+    ///             .body(ByteStream::from_static(format!("content for {}", req.key().unwrap_or("unknown")).as_bytes()))
+    ///             .build()
+    ///     })
+    ///     .build();
+    /// ```
+    pub fn then_compute_output<F>(self, compute_fn: F) -> Rule
+    where
+        F: Fn(&I) -> O + Send + Sync + 'static,
+    {
+        self.sequence().compute_output(compute_fn).build_simple()
+    }
 }
 
-type SequenceGeneratorFn<O, E> = Arc<dyn Fn() -> MockResponse<O, E> + Send + Sync>;
+type SequenceGeneratorFn<O, E> = Arc<dyn Fn(&Input) -> MockResponse<O, E> + Send + Sync>;
 
 /// A builder for creating response sequences
 pub struct ResponseSequenceBuilder<I, O, E> {
@@ -281,7 +304,7 @@ where
     where
         F: Fn() -> O + Send + Sync + 'static,
     {
-        let generator = Arc::new(move || MockResponse::Output(output_fn()));
+        let generator = Arc::new(move |_input: &Input| MockResponse::Output(output_fn()));
         self.generators.push((generator, 1));
         self
     }
@@ -291,7 +314,7 @@ where
     where
         F: Fn() -> E + Send + Sync + 'static,
     {
-        let generator = Arc::new(move || MockResponse::Error(error_fn()));
+        let generator = Arc::new(move |_input: &Input| MockResponse::Error(error_fn()));
         self.generators.push((generator, 1));
         self
     }
@@ -301,10 +324,10 @@ where
         let status_code = StatusCode::try_from(status).unwrap();
 
         let generator: SequenceGeneratorFn<O, E> = match body {
-            Some(body) => Arc::new(move || {
+            Some(body) => Arc::new(move |_input: &Input| {
                 MockResponse::Http(HttpResponse::new(status_code, SdkBody::from(body.clone())))
             }),
-            None => Arc::new(move || {
+            None => Arc::new(move |_input: &Input| {
                 MockResponse::Http(HttpResponse::new(status_code, SdkBody::empty()))
             }),
         };
@@ -318,7 +341,26 @@ where
     where
         F: Fn() -> HttpResponse + Send + Sync + 'static,
     {
-        let generator = Arc::new(move || MockResponse::Http(response_fn()));
+        let generator = Arc::new(move |_input: &Input| MockResponse::Http(response_fn()));
+        self.generators.push((generator, 1));
+        self
+    }
+
+    /// Add a computed output response to the sequence.  Note that this is not `pub`
+    /// because creating computed output rules off of sequenced rules doesn't work,
+    /// as we can't preserve the input across retries.  So we only expose `compute_output`
+    /// on unsequenced rules above.
+    fn compute_output<F>(mut self, compute_fn: F) -> Self
+    where
+        F: Fn(&I) -> O + Send + Sync + 'static,
+    {
+        let generator = Arc::new(move |input: &Input| {
+            if let Some(typed_input) = input.downcast_ref::<I>() {
+                MockResponse::Output(compute_fn(typed_input))
+            } else {
+                panic!("Input type mismatch in compute_output")
+            }
+        });
         self.generators.push((generator, 1));
         self
     }
@@ -415,12 +457,12 @@ where
 
         Rule::new(
             self.input_filter,
-            Arc::new(move |idx| {
+            Arc::new(move |idx, input| {
                 // find which generator to use
                 let mut current_idx = idx;
                 for (generator, repeat_count) in &generators {
                     if current_idx < *repeat_count {
-                        return Some(generator());
+                        return Some(generator(input));
                     }
                     current_idx -= repeat_count;
                 }