Add then_compute_response method to aws-smithy-mocks (#4388)

Adds a new method that allows developers to stub arbitrary responses
(outputs, errors, or HTTP responses) based on input request properties.

Unlike then_compute_output which only returns successful outputs, this
method enables conditional response generation where the response type
can vary based on the input.

Changes:
- Make MockResponse enum public to allow user construction
- Add RuleBuilder::then_compute_response() method
- Add ResponseSequenceBuilder::compute_response() internal method
- Export MockResponse from lib.rs
- Add integration test demonstrating conditional responses
- Update README with usage example

Example usage:
```
  mock!(Client::get_object) .then_compute_response(|req| { if req.key() == Some("error-key") { MockResponse::Error(...) } else { MockResponse::Output(...) } })
```

## Motivation and Context
Allow for stubbing/faking that sends failures back to the caller

## Testing
crate tests

## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [ ] For changes to the smithy-rs codegen or runtime crates, I have
created a changelog entry Markdown file in the `.changelog` directory,
specifying "client," "server," or both in the `applies_to` key.
- [ ] For changes to the AWS SDK, generated SDK code, or SDK runtime
crates, I have created a changelog entry Markdown file in the
`.changelog` directory, specifying "aws-sdk-rust" in the `applies_to`
key.

----

_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/1762538321.md

@@ -0,0 +1,11 @@
+---
+applies_to: ["client"]
+authors:
+- greenwoodcm
+references:
+- smithy-rs#4388
+breaking: false
+new_feature: false
+bug_fix: false
+---
+Add `then_compute_response` to Smithy mock

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

@@ -1,6 +1,6 @@
 [package]
 name = "aws-smithy-mocks"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>"]
 description = "Testing utilities for smithy-rs generated clients"
 edition = "2021"

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

@@ -117,6 +117,19 @@ let compute_rule = mock!(Client::get_object)
             .body(ByteStream::from_static(format!("content for {}", key).as_bytes()))
             .build()
     });
+
+// Return any response type (output, error, or HTTP) based on the input
+let conditional_rule = mock!(Client::get_object)
+    .then_compute_response(|req| {
+        use aws_smithy_mocks::MockResponse;
+        if req.key() == Some("error-key") {
+            MockResponse::Error(GetObjectError::NoSuchKey(NoSuchKey::builder().build()))
+        } else {
+            MockResponse::Output(GetObjectOutput::builder()
+                .body(ByteStream::from_static(b"success"))
+                .build())
+        }
+    });
 ```
 
 ### Response Sequences

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

@@ -911,4 +911,34 @@ mod tests {
             .build();
         assert_eq!(rule.max_responses, usize::MAX);
     }
+
+    #[tokio::test]
+    async fn test_compute_response_conditional() {
+        use crate::MockResponse;
+
+        let rule = create_rule_builder().then_compute_response(|input| {
+            if input.key == "error-key" {
+                MockResponse::Error(TestError::new("conditional error"))
+            } else {
+                MockResponse::Output(TestOutput::new(&format!("response for {}", input.key)))
+            }
+        });
+
+        let interceptor = MockResponseInterceptor::new().with_rule(&rule);
+        let operation = create_test_operation(interceptor, false);
+
+        // Test success case
+        let result = operation
+            .invoke(TestInput::new("test-bucket", "success-key"))
+            .await;
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), TestOutput::new("response for success-key"));
+
+        // Test error case
+        let result = operation
+            .invoke(TestInput::new("test-bucket", "error-key"))
+            .await;
+        assert!(result.is_err());
+        assert_eq!(rule.num_calls(), 2);
+    }
 }

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

@@ -18,8 +18,7 @@ mod interceptor;
 mod rule;
 
 pub use interceptor::{create_mock_http_client, MockResponseInterceptor};
-pub(crate) use rule::MockResponse;
-pub use rule::{Rule, RuleBuilder, RuleMode};
+pub use rule::{MockResponse, Rule, RuleBuilder, RuleMode};
 
 // why do we need a macro for this?
 // We want customers to be able to provide an ergonomic way to say the method they're looking for,

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

@@ -22,7 +22,7 @@ use std::sync::Arc;
 ///
 #[derive(Debug)]
 #[allow(clippy::large_enum_variant)]
-pub(crate) enum MockResponse<O, E> {
+pub enum MockResponse<O, E> {
     /// A successful modeled response.
     Output(O),
     /// A modeled error.
@@ -251,6 +251,33 @@ where
     {
         self.sequence().compute_output(compute_fn).build_simple()
     }
+
+    /// Creates a rule that computes an arbitrary response based on the input.
+    ///
+    /// This allows generating any type of response (output, error, or HTTP) based on the input request.
+    /// Unlike `then_compute_output`, this method can return errors or HTTP responses conditionally.
+    ///
+    /// # Examples
+    ///
+    /// ```rust,ignore
+    /// let rule = mock!(Client::get_object)
+    ///     .then_compute_response(|req| {
+    ///         if req.key() == Some("error") {
+    ///             MockResponse::Error(GetObjectError::NoSuchKey(NoSuchKey::builder().build()))
+    ///         } else {
+    ///             MockResponse::Output(GetObjectOutput::builder()
+    ///                 .body(ByteStream::from_static(b"content"))
+    ///                 .build())
+    ///         }
+    ///     })
+    ///     .build();
+    /// ```
+    pub fn then_compute_response<F>(self, compute_fn: F) -> Rule
+    where
+        F: Fn(&I) -> MockResponse<O, E> + Send + Sync + 'static,
+    {
+        self.sequence().compute_response(compute_fn).build_simple()
+    }
 }
 
 type SequenceGeneratorFn<O, E> = Arc<dyn Fn(&Input) -> MockResponse<O, E> + Send + Sync>;
@@ -366,6 +393,22 @@ where
         self
     }
 
+    /// Add a computed response to the sequence. Not `pub` for same reason as `compute_output`.
+    fn compute_response<F>(mut self, compute_fn: F) -> Self
+    where
+        F: Fn(&I) -> MockResponse<O, E> + Send + Sync + 'static,
+    {
+        let generator = Arc::new(move |input: &Input| {
+            if let Some(typed_input) = input.downcast_ref::<I>() {
+                compute_fn(typed_input)
+            } else {
+                panic!("Input type mismatch in compute_response")
+            }
+        });
+        self.generators.push((generator, 1));
+        self
+    }
+
     /// Repeat the last added response multiple times.
     ///
     /// This method sets the number of times the last response in the sequence will be used.