aws_smithy_mocks/
interceptor.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use crate::{MockResponse, Rule, RuleMode};
7use aws_smithy_http_client::test_util::infallible_client_fn;
8use aws_smithy_runtime_api::box_error::BoxError;
9use aws_smithy_runtime_api::client::http::SharedHttpClient;
10use aws_smithy_runtime_api::client::interceptors::context::{
11    BeforeSerializationInterceptorContextMut, BeforeTransmitInterceptorContextMut, Error,
12    FinalizerInterceptorContextMut, Input, Output,
13};
14use aws_smithy_runtime_api::client::interceptors::Intercept;
15use aws_smithy_runtime_api::client::orchestrator::{HttpResponse, OrchestratorError};
16use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
17use aws_smithy_types::body::SdkBody;
18use aws_smithy_types::config_bag::{ConfigBag, Storable, StoreReplace};
19use std::collections::VecDeque;
20use std::fmt;
21use std::sync::{Arc, Mutex};
22
23// Store active rule in config bag
24#[derive(Debug, Clone)]
25struct ActiveRule(Rule);
26
27impl Storable for ActiveRule {
28    type Storer = StoreReplace<ActiveRule>;
29}
30
31/// Interceptor which produces mock responses based on a list of rules
32pub struct MockResponseInterceptor {
33    rules: Arc<Mutex<VecDeque<Rule>>>,
34    rule_mode: RuleMode,
35    must_match: bool,
36    active_response: Arc<Mutex<Option<MockResponse<Output, Error>>>>,
37}
38
39impl fmt::Debug for MockResponseInterceptor {
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        write!(f, "{} rules", self.rules.lock().unwrap().len())
42    }
43}
44
45impl Default for MockResponseInterceptor {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl MockResponseInterceptor {
52    /// Create a new [MockResponseInterceptor]
53    ///
54    /// This is normally created and registered on a client through the [`mock_client`](crate::mock_client) macro.
55    pub fn new() -> Self {
56        Self {
57            rules: Default::default(),
58            rule_mode: RuleMode::MatchAny,
59            must_match: true,
60            active_response: Default::default(),
61        }
62    }
63    /// Add a rule to the Interceptor
64    ///
65    /// Rules are matched in order—this rule will only apply if all previous rules do not match.
66    pub fn with_rule(self, rule: &Rule) -> Self {
67        self.rules.lock().unwrap().push_back(rule.clone());
68        self
69    }
70
71    /// Set the RuleMode to use when evaluating rules.
72    ///
73    /// See `RuleMode` enum for modes and how they are applied.
74    pub fn rule_mode(mut self, rule_mode: RuleMode) -> Self {
75        self.rule_mode = rule_mode;
76        self
77    }
78
79    /// Allow passthrough for unmatched requests.
80    ///
81    /// By default, if a request doesn't match any rule, the interceptor will panic.
82    /// This method allows unmatched requests to pass through.
83    pub fn allow_passthrough(mut self) -> Self {
84        self.must_match = false;
85        self
86    }
87}
88
89impl Intercept for MockResponseInterceptor {
90    fn name(&self) -> &'static str {
91        "MockResponseInterceptor"
92    }
93
94    fn modify_before_serialization(
95        &self,
96        context: &mut BeforeSerializationInterceptorContextMut<'_>,
97        _runtime_components: &RuntimeComponents,
98        cfg: &mut ConfigBag,
99    ) -> Result<(), BoxError> {
100        let mut rules = self.rules.lock().unwrap();
101        let input = context.inner().input().expect("input set");
102
103        // Find a matching rule and get its response
104        let mut matching_rule = None;
105        let mut matching_response = None;
106
107        match self.rule_mode {
108            RuleMode::Sequential => {
109                // Sequential mode requires rules match in-order
110                let i = 0;
111                while i < rules.len() && matching_response.is_none() {
112                    let rule = &rules[i];
113
114                    // Check if the rule is already exhausted or if it's a simple rule used once
115                    //
116                    // In `aws-smithy-mocks-experimental` all rules were infinite sequences
117                    // but were only usable once in sequential mode. We retain that here for
118                    // backwards compatibility.
119                    if rule.is_exhausted() || (rule.is_simple() && rule.num_calls() > 0) {
120                        // Rule is exhausted, remove it and try the next one
121                        rules.remove(i);
122                        continue; // Don't increment i since we removed an element
123                    }
124
125                    // Check if the rule matches
126                    if !(rule.matcher)(input) {
127                        // Rule doesn't match, this is an error in sequential mode
128                        panic!(
129                            "In order matching was enforced but rule did not match {:?}",
130                            input
131                        );
132                    }
133
134                    // Rule matches and is not exhausted, get the response
135                    if let Some(response) = rule.next_response(input) {
136                        matching_rule = Some(rule.clone());
137                        matching_response = Some(response);
138                    } else {
139                        // Rule is exhausted, remove it and try the next one
140                        rules.remove(i);
141                        continue; // Don't increment i since we removed an element
142                    }
143
144                    // We found a matching rule and got a response, so we're done
145                    break;
146                }
147            }
148            RuleMode::MatchAny => {
149                // Find any matching rule with a response
150                for rule in rules.iter() {
151                    // Skip exhausted rules
152                    if rule.is_exhausted() {
153                        continue;
154                    }
155
156                    if (rule.matcher)(input) {
157                        if let Some(response) = rule.next_response(input) {
158                            matching_rule = Some(rule.clone());
159                            matching_response = Some(response);
160                            break;
161                        }
162                    }
163                }
164            }
165        };
166
167        match (matching_rule, matching_response) {
168            (Some(rule), Some(response)) => {
169                // Store the rule in the config bag
170                cfg.interceptor_state().store_put(ActiveRule(rule));
171                // store the response on the interceptor (because going
172                // through interceptor context requires the type to impl Clone)
173                let mut active_resp = self.active_response.lock().unwrap();
174                let _ = std::mem::replace(&mut *active_resp, Some(response));
175            }
176            _ => {
177                // No matching rule or no response
178                if self.must_match {
179                    panic!(
180                        "must_match was enabled but no rules matched or all rules were exhausted for {:?}",
181                        input
182                    );
183                }
184            }
185        }
186
187        Ok(())
188    }
189
190    fn modify_before_transmit(
191        &self,
192        context: &mut BeforeTransmitInterceptorContextMut<'_>,
193        _runtime_components: &RuntimeComponents,
194        cfg: &mut ConfigBag,
195    ) -> Result<(), BoxError> {
196        let mut state = self.active_response.lock().unwrap();
197        let mut active_response = (*state).take();
198        if active_response.is_none() {
199            // in the case of retries we try to get the next response if it has been consumed
200            if let Some(active_rule) = cfg.load::<ActiveRule>() {
201                // During retries, input is not available in modify_before_transmit.
202                // For HTTP status responses that don't use the input, we can use a dummy input.
203                let dummy_input = Input::doesnt_matter();
204                let next_resp = active_rule.0.next_response(&dummy_input);
205                active_response = next_resp;
206            }
207        }
208
209        if let Some(resp) = active_response {
210            match resp {
211                // place the http response into the extensions and let the HTTP client return it
212                MockResponse::Http(http_resp) => {
213                    context
214                        .request_mut()
215                        .add_extension(MockHttpResponse(Arc::new(http_resp)));
216                }
217                _ => {
218                    // put it back for modeled output/errors
219                    let _ = std::mem::replace(&mut *state, Some(resp));
220                }
221            }
222        }
223
224        Ok(())
225    }
226
227    fn modify_before_attempt_completion(
228        &self,
229        context: &mut FinalizerInterceptorContextMut<'_>,
230        _runtime_components: &RuntimeComponents,
231        _cfg: &mut ConfigBag,
232    ) -> Result<(), BoxError> {
233        // Handle modeled responses
234        let mut state = self.active_response.lock().unwrap();
235        let active_response = (*state).take();
236        if let Some(resp) = active_response {
237            match resp {
238                MockResponse::Output(output) => {
239                    context.inner_mut().set_output_or_error(Ok(output));
240                }
241                MockResponse::Error(error) => {
242                    context
243                        .inner_mut()
244                        .set_output_or_error(Err(OrchestratorError::operation(error)));
245                }
246                MockResponse::Http(_) => {
247                    // HTTP responses are handled by the mock HTTP client
248                }
249            }
250        }
251
252        Ok(())
253    }
254}
255
256/// Extension for storing mock HTTP responses in request extensions
257#[derive(Clone)]
258struct MockHttpResponse(Arc<HttpResponse>);
259
260/// Create a mock HTTP client that works with the interceptor using existing utilities
261pub fn create_mock_http_client() -> SharedHttpClient {
262    infallible_client_fn(|mut req| {
263        // Try to get the mock HTTP response generator from the extensions
264        if let Some(mock_response) = req.extensions_mut().remove::<MockHttpResponse>() {
265            let http_resp =
266                Arc::try_unwrap(mock_response.0).expect("mock HTTP response has single reference");
267            return http_resp.try_into_http1x().unwrap();
268        }
269
270        // Default dummy response if no mock response is defined
271        http::Response::builder()
272            .status(418)
273            .body(SdkBody::from("Mock HTTP client dummy response"))
274            .unwrap()
275    })
276}
277
278#[cfg(test)]
279mod tests {
280    use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep};
281    use aws_smithy_runtime::client::orchestrator::operation::Operation;
282    use aws_smithy_runtime::client::retries::classifiers::HttpStatusCodeClassifier;
283    use aws_smithy_runtime_api::client::orchestrator::{
284        HttpRequest, HttpResponse, OrchestratorError,
285    };
286    use aws_smithy_runtime_api::client::result::SdkError;
287    use aws_smithy_runtime_api::http::StatusCode;
288    use aws_smithy_types::body::SdkBody;
289    use aws_smithy_types::retry::RetryConfig;
290    use aws_smithy_types::timeout::TimeoutConfig;
291
292    use crate::{create_mock_http_client, MockResponseInterceptor, RuleBuilder, RuleMode};
293    use std::time::Duration;
294
295    // Simple test input and output types
296    #[derive(Debug)]
297    struct TestInput {
298        bucket: String,
299        key: String,
300    }
301    impl TestInput {
302        fn new(bucket: &str, key: &str) -> Self {
303            Self {
304                bucket: bucket.to_string(),
305                key: key.to_string(),
306            }
307        }
308    }
309
310    #[derive(Debug, PartialEq)]
311    struct TestOutput {
312        content: String,
313    }
314
315    impl TestOutput {
316        fn new(content: &str) -> Self {
317            Self {
318                content: content.to_string(),
319            }
320        }
321    }
322
323    #[derive(Debug)]
324    struct TestError {
325        message: String,
326    }
327
328    impl TestError {
329        fn new(message: &str) -> Self {
330            Self {
331                message: message.to_string(),
332            }
333        }
334    }
335
336    impl std::fmt::Display for TestError {
337        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
338            write!(f, "{}", self.message)
339        }
340    }
341
342    impl std::error::Error for TestError {}
343
344    // Helper function to create a RuleBuilder with proper type hints
345    fn create_rule_builder() -> RuleBuilder<TestInput, TestOutput, TestError> {
346        RuleBuilder::new_from_mock(
347            || TestInput {
348                bucket: "".to_string(),
349                key: "".to_string(),
350            },
351            || {
352                let fut: std::future::Ready<Result<TestOutput, SdkError<TestError, HttpResponse>>> =
353                    std::future::ready(Ok(TestOutput {
354                        content: "".to_string(),
355                    }));
356                fut
357            },
358        )
359    }
360
361    // Helper function to create an Operation with common configuration
362    fn create_test_operation(
363        interceptor: MockResponseInterceptor,
364        enable_retries: bool,
365    ) -> Operation<TestInput, TestOutput, TestError> {
366        let builder = Operation::builder()
367            .service_name("test")
368            .operation_name("test")
369            .http_client(create_mock_http_client())
370            .endpoint_url("http://localhost:1234")
371            .no_auth()
372            .sleep_impl(SharedAsyncSleep::new(TokioSleep::new()))
373            .timeout_config(TimeoutConfig::disabled())
374            .interceptor(interceptor)
375            .serializer(|input: TestInput| {
376                let mut request = HttpRequest::new(SdkBody::empty());
377                request
378                    .set_uri(format!("/{}/{}", input.bucket, input.key))
379                    .expect("valid URI");
380                Ok(request)
381            })
382            .deserializer::<TestOutput, TestError>(|response| {
383                if response.status().is_success() {
384                    let body = std::str::from_utf8(response.body().bytes().unwrap())
385                        .unwrap_or("empty body")
386                        .to_string();
387                    Ok(TestOutput { content: body })
388                } else {
389                    Err(OrchestratorError::operation(TestError {
390                        message: format!("Error: {}", response.status()),
391                    }))
392                }
393            });
394
395        if enable_retries {
396            let retry_config = RetryConfig::standard()
397                .with_max_attempts(5)
398                .with_initial_backoff(Duration::from_millis(1))
399                .with_max_backoff(Duration::from_millis(5));
400
401            builder
402                .retry_classifier(HttpStatusCodeClassifier::default())
403                .standard_retry(&retry_config)
404                .build()
405        } else {
406            builder.no_retry().build()
407        }
408    }
409
410    #[tokio::test]
411    async fn test_retry_sequence() {
412        // Create a rule with repeated error responses followed by success
413        let rule = create_rule_builder()
414            .match_requests(|input| input.bucket == "test-bucket" && input.key == "test-key")
415            .sequence()
416            .http_status(503, None)
417            .times(2)
418            .output(|| TestOutput::new("success after retries"))
419            .build();
420
421        // Create an interceptor with the rule
422        let interceptor = MockResponseInterceptor::new()
423            .rule_mode(RuleMode::Sequential)
424            .with_rule(&rule);
425
426        let operation = create_test_operation(interceptor, true);
427
428        // Make a single request - it should automatically retry through the sequence
429        let result = operation
430            .invoke(TestInput::new("test-bucket", "test-key"))
431            .await;
432
433        // Should succeed with the final output after retries
434        assert!(
435            result.is_ok(),
436            "Expected success but got error: {:?}",
437            result.err()
438        );
439        assert_eq!(
440            result.unwrap(),
441            TestOutput {
442                content: "success after retries".to_string()
443            }
444        );
445
446        // Verify the rule was used the expected number of times (all 4 responses: 2 errors + 1 success)
447        assert_eq!(rule.num_calls(), 3);
448    }
449
450    #[tokio::test]
451    async fn test_compute_output() {
452        // Create a rule that computes its responses based off of input data
453        let rule = create_rule_builder()
454            .match_requests(|input| input.bucket == "test-bucket" && input.key == "test-key")
455            .then_compute_output(|input| TestOutput {
456                content: format!("{}.{}", input.bucket, input.key),
457            });
458
459        // Create an interceptor with the rule
460        let interceptor = MockResponseInterceptor::new()
461            .rule_mode(RuleMode::Sequential)
462            .with_rule(&rule);
463
464        let operation = create_test_operation(interceptor, true);
465
466        let result = operation
467            .invoke(TestInput::new("test-bucket", "test-key"))
468            .await;
469
470        // Should succeed with the output derived from input
471        assert!(
472            result.is_ok(),
473            "Expected success but got error: {:?}",
474            result.err()
475        );
476        assert_eq!(
477            result.unwrap(),
478            TestOutput {
479                content: "test-bucket.test-key".to_string()
480            }
481        );
482
483        // Verify the rule was used once, no retries
484        assert_eq!(rule.num_calls(), 1);
485    }
486
487    #[should_panic(
488        expected = "must_match was enabled but no rules matched or all rules were exhausted for"
489    )]
490    #[tokio::test]
491    async fn test_exhausted_rules_sequential() {
492        // Create a rule with a single response
493        let rule = create_rule_builder().then_output(|| TestOutput::new("only response"));
494
495        // Create an interceptor with the rule
496        let interceptor = MockResponseInterceptor::new()
497            .rule_mode(RuleMode::Sequential)
498            .with_rule(&rule);
499
500        let operation = create_test_operation(interceptor, false);
501
502        // First call should succeed
503        let result1 = operation
504            .invoke(TestInput::new("test-bucket", "test-key"))
505            .await;
506        assert!(result1.is_ok());
507
508        // Second call should panic because the rules are exhausted
509        let _result2 = operation
510            .invoke(TestInput::new("test-bucket", "test-key"))
511            .await;
512    }
513
514    #[tokio::test]
515    async fn test_rule_mode_match_any() {
516        // Create two rules with different matchers
517        let rule1 = create_rule_builder()
518            .match_requests(|input| input.bucket == "bucket1")
519            .then_output(|| TestOutput::new("response1"));
520
521        let rule2 = create_rule_builder()
522            .match_requests(|input| input.bucket == "bucket2")
523            .then_output(|| TestOutput::new("response2"));
524
525        // Create an interceptor with both rules in MatchAny mode
526        let interceptor = MockResponseInterceptor::new()
527            .rule_mode(RuleMode::MatchAny)
528            .with_rule(&rule1)
529            .with_rule(&rule2);
530
531        let operation = create_test_operation(interceptor, false);
532
533        // Call with bucket1 should match rule1
534        let result1 = operation
535            .invoke(TestInput::new("bucket1", "test-key"))
536            .await;
537        assert!(result1.is_ok());
538        assert_eq!(result1.unwrap(), TestOutput::new("response1"));
539
540        // Call with bucket2 should match rule2
541        let result2 = operation
542            .invoke(TestInput::new("bucket2", "test-key"))
543            .await;
544        assert!(result2.is_ok());
545        assert_eq!(result2.unwrap(), TestOutput::new("response2"));
546
547        // Verify the rules were used the expected number of times
548        assert_eq!(rule1.num_calls(), 1);
549        assert_eq!(rule2.num_calls(), 1);
550
551        // Calling with bucket1 again should match rule1 a second time
552        let result1 = operation
553            .invoke(TestInput::new("bucket1", "test-key"))
554            .await;
555        assert!(result1.is_ok());
556        assert_eq!(result1.unwrap(), TestOutput::new("response1"));
557        assert_eq!(rule1.num_calls(), 2);
558    }
559
560    #[tokio::test]
561    async fn test_mixed_response_types() {
562        // Create a rule with all three types of responses
563        let rule = create_rule_builder()
564            .sequence()
565            .output(|| TestOutput::new("first output"))
566            .error(|| TestError::new("expected error"))
567            .http_response(|| {
568                HttpResponse::new(
569                    StatusCode::try_from(200).unwrap(),
570                    SdkBody::from("http response"),
571                )
572            })
573            .build();
574
575        // Create an interceptor with the rule
576        let interceptor = MockResponseInterceptor::new()
577            .rule_mode(RuleMode::Sequential)
578            .with_rule(&rule);
579
580        let operation = create_test_operation(interceptor, false);
581
582        // First call should return the modeled output
583        let result1 = operation
584            .invoke(TestInput::new("test-bucket", "test-key"))
585            .await;
586        assert!(result1.is_ok());
587        assert_eq!(result1.unwrap(), TestOutput::new("first output"));
588
589        // Second call should return the modeled error
590        let result2 = operation
591            .invoke(TestInput::new("test-bucket", "test-key"))
592            .await;
593        assert!(result2.is_err());
594        let sdk_err = result2.unwrap_err();
595        let err = sdk_err.as_service_error().expect("expected service error");
596        assert_eq!(err.to_string(), "expected error");
597
598        // Third call should return the HTTP response
599        let result3 = operation
600            .invoke(TestInput::new("test-bucket", "test-key"))
601            .await;
602        assert!(result3.is_ok());
603        assert_eq!(result3.unwrap(), TestOutput::new("http response"));
604
605        // Verify the rule was used the expected number of times
606        assert_eq!(rule.num_calls(), 3);
607    }
608    #[tokio::test]
609    async fn test_exhausted_sequence_match_any() {
610        // Create a rule with a sequence that will be exhausted
611        let rule = create_rule_builder()
612            .match_requests(|input| input.bucket == "bucket-1")
613            .sequence()
614            .output(|| TestOutput::new("response 1"))
615            .output(|| TestOutput::new("response 2"))
616            .build();
617
618        // Create another rule to use after the first one is exhausted
619        let fallback_rule =
620            create_rule_builder().then_output(|| TestOutput::new("fallback response"));
621
622        // Create an interceptor with both rules
623        let interceptor = MockResponseInterceptor::new()
624            .rule_mode(RuleMode::MatchAny)
625            .with_rule(&rule)
626            .with_rule(&fallback_rule);
627
628        let operation = create_test_operation(interceptor, false);
629
630        // First two calls should use the first rule
631        let result1 = operation
632            .invoke(TestInput::new("bucket-1", "test-key"))
633            .await;
634        assert!(result1.is_ok());
635        assert_eq!(result1.unwrap(), TestOutput::new("response 1"));
636
637        // second should use our fallback rule
638        let result2 = operation
639            .invoke(TestInput::new("other-bucket", "test-key"))
640            .await;
641        assert!(result2.is_ok());
642        assert_eq!(result2.unwrap(), TestOutput::new("fallback response"));
643
644        // Third call should use the first rule again and exhaust it
645        let result3 = operation
646            .invoke(TestInput::new("bucket-1", "test-key"))
647            .await;
648        assert!(result3.is_ok());
649        assert_eq!(result3.unwrap(), TestOutput::new("response 2"));
650
651        // first rule is exhausted so the matcher shouldn't matter and we should hit our fallback rule
652        let result4 = operation
653            .invoke(TestInput::new("bucket-1", "test-key"))
654            .await;
655        assert!(result4.is_ok());
656        assert_eq!(result4.unwrap(), TestOutput::new("fallback response"));
657
658        // Verify the rules were used the expected number of times
659        assert_eq!(rule.num_calls(), 2);
660        assert_eq!(fallback_rule.num_calls(), 2);
661    }
662
663    #[tokio::test]
664    async fn test_exhausted_sequence_sequential() {
665        // Create a rule with a sequence that will be exhausted
666        let rule = create_rule_builder()
667            .sequence()
668            .output(|| TestOutput::new("response 1"))
669            .output(|| TestOutput::new("response 2"))
670            .build();
671
672        // Create another rule to use after the first one is exhausted
673        let fallback_rule =
674            create_rule_builder().then_output(|| TestOutput::new("fallback response"));
675
676        // Create an interceptor with both rules
677        let interceptor = MockResponseInterceptor::new()
678            .rule_mode(RuleMode::Sequential)
679            .with_rule(&rule)
680            .with_rule(&fallback_rule);
681
682        let operation = create_test_operation(interceptor, false);
683
684        // First two calls should use the first rule
685        let result1 = operation
686            .invoke(TestInput::new("test-bucket", "test-key"))
687            .await;
688        assert!(result1.is_ok());
689        assert_eq!(result1.unwrap(), TestOutput::new("response 1"));
690
691        let result2 = operation
692            .invoke(TestInput::new("test-bucket", "test-key"))
693            .await;
694        assert!(result2.is_ok());
695        assert_eq!(result2.unwrap(), TestOutput::new("response 2"));
696
697        // Third call should use the fallback rule
698        let result3 = operation
699            .invoke(TestInput::new("test-bucket", "test-key"))
700            .await;
701        assert!(result3.is_ok());
702        assert_eq!(result3.unwrap(), TestOutput::new("fallback response"));
703
704        // Verify the rules were used the expected number of times
705        assert_eq!(rule.num_calls(), 2);
706        assert_eq!(fallback_rule.num_calls(), 1);
707    }
708
709    #[tokio::test]
710    async fn test_concurrent_usage() {
711        use std::sync::Arc;
712        use tokio::task;
713
714        // Create a rule with multiple responses
715        let rule = Arc::new(
716            create_rule_builder()
717                .sequence()
718                .output(|| TestOutput::new("response 1"))
719                .output(|| TestOutput::new("response 2"))
720                .output(|| TestOutput::new("response 3"))
721                .build(),
722        );
723
724        // Create an interceptor with the rule
725        let interceptor = MockResponseInterceptor::new()
726            .rule_mode(RuleMode::Sequential)
727            .with_rule(&rule);
728
729        let operation = Arc::new(create_test_operation(interceptor, false));
730
731        // Spawn multiple tasks that use the operation concurrently
732        let mut handles = vec![];
733        for i in 0..3 {
734            let op = operation.clone();
735            let handle = task::spawn(async move {
736                let result = op
737                    .invoke(TestInput::new(&format!("bucket-{}", i), "test-key"))
738                    .await;
739                result.unwrap()
740            });
741            handles.push(handle);
742        }
743
744        // Wait for all tasks to complete
745        let mut results = vec![];
746        for handle in handles {
747            results.push(handle.await.unwrap());
748        }
749
750        // Sort the results to make the test deterministic
751        results.sort_by(|a, b| a.content.cmp(&b.content));
752
753        // Verify we got all three responses
754        assert_eq!(results.len(), 3);
755        assert_eq!(results[0], TestOutput::new("response 1"));
756        assert_eq!(results[1], TestOutput::new("response 2"));
757        assert_eq!(results[2], TestOutput::new("response 3"));
758
759        // Verify the rule was used the expected number of times
760        assert_eq!(rule.num_calls(), 3);
761    }
762
763    #[tokio::test]
764    async fn test_sequential_rule_removal() {
765        // Create a rule that matches only when key != "correct-key"
766        let rule1 = create_rule_builder()
767            .match_requests(|input| input.bucket == "test-bucket" && input.key != "correct-key")
768            .then_http_response(|| {
769                HttpResponse::new(
770                    StatusCode::try_from(404).unwrap(),
771                    SdkBody::from("not found"),
772                )
773            });
774
775        // Create a rule that matches only when key == "correct-key"
776        let rule2 = create_rule_builder()
777            .match_requests(|input| input.bucket == "test-bucket" && input.key == "correct-key")
778            .then_output(|| TestOutput::new("success"));
779
780        // Create an interceptor with both rules in Sequential mode
781        let interceptor = MockResponseInterceptor::new()
782            .rule_mode(RuleMode::Sequential)
783            .with_rule(&rule1)
784            .with_rule(&rule2);
785
786        let operation = create_test_operation(interceptor, true);
787
788        // First call with key="foo" should match rule1
789        let result1 = operation.invoke(TestInput::new("test-bucket", "foo")).await;
790        assert!(result1.is_err());
791        assert_eq!(rule1.num_calls(), 1);
792
793        // Second call with key="correct-key" should match rule2
794        // But this will fail if rule1 is not removed after being used
795        let result2 = operation
796            .invoke(TestInput::new("test-bucket", "correct-key"))
797            .await;
798
799        // This should succeed, rule1 doesn't match but should have been removed
800        assert!(result2.is_ok());
801        assert_eq!(result2.unwrap(), TestOutput::new("success"));
802        assert_eq!(rule2.num_calls(), 1);
803    }
804
805    #[tokio::test]
806    async fn test_simple_rule_in_match_any_mode() {
807        let rule = create_rule_builder().then_output(|| TestOutput::new("simple response"));
808
809        let interceptor = MockResponseInterceptor::new()
810            .rule_mode(RuleMode::MatchAny)
811            .with_rule(&rule);
812
813        let operation = create_test_operation(interceptor, false);
814
815        for i in 0..5 {
816            let result = operation
817                .invoke(TestInput::new("test-bucket", "test-key"))
818                .await;
819            assert!(result.is_ok(), "Call {} should succeed", i);
820            assert_eq!(result.unwrap(), TestOutput::new("simple response"));
821        }
822        assert_eq!(rule.num_calls(), 5);
823        assert!(!rule.is_exhausted());
824    }
825
826    #[tokio::test]
827    async fn test_simple_rule_in_sequential_mode() {
828        let rule1 = create_rule_builder().then_output(|| TestOutput::new("first response"));
829        let rule2 = create_rule_builder().then_output(|| TestOutput::new("second response"));
830
831        let interceptor = MockResponseInterceptor::new()
832            .rule_mode(RuleMode::Sequential)
833            .with_rule(&rule1)
834            .with_rule(&rule2);
835
836        let operation = create_test_operation(interceptor, false);
837
838        let result1 = operation
839            .invoke(TestInput::new("test-bucket", "test-key"))
840            .await;
841        assert!(result1.is_ok());
842        assert_eq!(result1.unwrap(), TestOutput::new("first response"));
843
844        // Second call should use rule2 (rule1 should be removed after one use in Sequential mode)
845        let result2 = operation
846            .invoke(TestInput::new("test-bucket", "test-key"))
847            .await;
848        assert!(result2.is_ok());
849        assert_eq!(result2.unwrap(), TestOutput::new("second response"));
850
851        assert_eq!(rule1.num_calls(), 1);
852        assert_eq!(rule2.num_calls(), 1);
853    }
854
855    #[tokio::test]
856    async fn test_repeatedly_method() {
857        let rule = create_rule_builder()
858            .sequence()
859            .output(|| TestOutput::new("first response"))
860            .output(|| TestOutput::new("repeated response"))
861            .repeatedly()
862            .build();
863
864        let interceptor = MockResponseInterceptor::new()
865            .rule_mode(RuleMode::Sequential)
866            .with_rule(&rule);
867
868        let operation = create_test_operation(interceptor, false);
869
870        let result1 = operation
871            .invoke(TestInput::new("test-bucket", "test-key"))
872            .await;
873        assert!(result1.is_ok());
874        assert_eq!(result1.unwrap(), TestOutput::new("first response"));
875
876        // all subsequent calls should return "repeated response"
877        for i in 0..10 {
878            let result = operation
879                .invoke(TestInput::new("test-bucket", "test-key"))
880                .await;
881            assert!(result.is_ok(), "Call {} should succeed", i);
882            assert_eq!(result.unwrap(), TestOutput::new("repeated response"));
883        }
884        assert_eq!(rule.num_calls(), 11);
885        assert!(!rule.is_exhausted());
886    }
887
888    #[should_panic(expected = "times(n) called before adding a response to the sequence")]
889    #[test]
890    fn test_times_validation() {
891        // This should panic because times() is called before adding any responses
892        let _rule = create_rule_builder()
893            .sequence()
894            .times(3)
895            .output(|| TestOutput::new("response"))
896            .build();
897    }
898
899    #[should_panic(expected = "repeatedly() called before adding a response to the sequence")]
900    #[test]
901    fn test_repeatedly_validation() {
902        // This should panic because repeatedly() is called before adding any responses
903        let _rule = create_rule_builder().sequence().repeatedly().build();
904    }
905
906    #[test]
907    fn test_total_responses_overflow() {
908        // Create a rule with a large number of repetitions to test overflow handling
909        let rule = create_rule_builder()
910            .sequence()
911            .output(|| TestOutput::new("response"))
912            .times(usize::MAX / 2)
913            .output(|| TestOutput::new("another response"))
914            .repeatedly()
915            .build();
916        assert_eq!(rule.max_responses, usize::MAX);
917    }
918}