From a27be2bad80b916e431968b315cfa67cc42ac388 Mon Sep 17 00:00:00 2001
From: Russell Cohen <rcoh@amazon.com>
Date: Wed, 20 Dec 2023 10:25:29 -0500
Subject: [PATCH] Add dynamodb retry test (#3321)

## Motivation and Context
Adds a test to ensure that `ThrottlingErrors` from DynamoDB do not
result in a new connection


## Checklist
- no changelog

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
---
 aws/sdk/integration-tests/dynamodb/Cargo.toml |  2 +-
 .../tests/test-error-classification.rs        | 70 +++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 aws/sdk/integration-tests/dynamodb/tests/test-error-classification.rs

diff --git a/aws/sdk/integration-tests/dynamodb/Cargo.toml b/aws/sdk/integration-tests/dynamodb/Cargo.toml
index f1611ef47..57eadecc0 100644
--- a/aws/sdk/integration-tests/dynamodb/Cargo.toml
+++ b/aws/sdk/integration-tests/dynamodb/Cargo.toml
@@ -19,7 +19,7 @@ aws-sdk-dynamodb = { path = "../../build/aws-sdk/sdk/dynamodb", features = ["beh
 aws-smithy-async = { path = "../../build/aws-sdk/sdk/aws-smithy-async", features = ["test-util"] }
 aws-smithy-http = { path = "../../build/aws-sdk/sdk/aws-smithy-http" }
 aws-smithy-protocol-test = { path = "../../build/aws-sdk/sdk/aws-smithy-protocol-test" }
-aws-smithy-runtime = { path = "../../build/aws-sdk/sdk/aws-smithy-runtime", features = ["test-util"]}
+aws-smithy-runtime = { path = "../../build/aws-sdk/sdk/aws-smithy-runtime", features = ["test-util", "wire-mock"]}
 aws-smithy-runtime-api = { path = "../../build/aws-sdk/sdk/aws-smithy-runtime-api", features = ["test-util"]}
 aws-smithy-types = { path = "../../build/aws-sdk/sdk/aws-smithy-types", features = ["test-util"]}
 aws-types = { path = "../../build/aws-sdk/sdk/aws-types" }
diff --git a/aws/sdk/integration-tests/dynamodb/tests/test-error-classification.rs b/aws/sdk/integration-tests/dynamodb/tests/test-error-classification.rs
new file mode 100644
index 000000000..244f20c87
--- /dev/null
+++ b/aws/sdk/integration-tests/dynamodb/tests/test-error-classification.rs
@@ -0,0 +1,70 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use std::time::Duration;
+
+use aws_credential_types::Credentials;
+use aws_sdk_dynamodb::types::AttributeValue;
+use aws_sdk_dynamodb::Client;
+use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep};
+use aws_smithy_runtime::client::http::test_util::wire::{ReplayedEvent, WireMockServer};
+use aws_smithy_runtime::{ev, match_events};
+use aws_smithy_types::retry::RetryConfig;
+use aws_smithy_types::timeout::TimeoutConfig;
+use aws_types::region::Region;
+use bytes::Bytes;
+
+const DYNAMO_THROTTLING_RESPONSE: &str = r#"{"__type":"com.amazonaws.dynamodb.v20120810#ThrottlingException",
+"message":"enhance your calm"}"#;
+
+const DYNAMODB_DB_SUCCESS_RESPONSE: &str = r#"{"Count":0,"Items":[],"ScannedCount":2}"#;
+
+#[tokio::test]
+async fn test_no_reconnect_500_throttling() {
+    assert_error_not_transient(ReplayedEvent::HttpResponse {
+        status: 500,
+        body: Bytes::from(DYNAMO_THROTTLING_RESPONSE),
+    })
+    .await
+}
+
+#[tokio::test]
+async fn test_no_reconnect_429_throttling() {
+    assert_error_not_transient(ReplayedEvent::HttpResponse {
+        status: 429,
+        body: Bytes::from(DYNAMO_THROTTLING_RESPONSE),
+    })
+    .await
+}
+
+async fn assert_error_not_transient(error: ReplayedEvent) {
+    let mock = WireMockServer::start(vec![
+        error,
+        ReplayedEvent::with_body(DYNAMODB_DB_SUCCESS_RESPONSE),
+    ])
+    .await;
+
+    let config = aws_sdk_dynamodb::Config::builder()
+        .region(Region::from_static("us-east-2"))
+        .credentials_provider(Credentials::for_tests())
+        .sleep_impl(SharedAsyncSleep::new(TokioSleep::new()))
+        .endpoint_url(mock.endpoint_url())
+        .http_client(mock.http_client())
+        .timeout_config(
+            TimeoutConfig::builder()
+                .operation_attempt_timeout(Duration::from_millis(100))
+                .build(),
+        )
+        .retry_config(RetryConfig::standard())
+        .build();
+    let client = Client::from_conf(config);
+    let _item = client
+        .get_item()
+        .key("foo", AttributeValue::Bool(true))
+        .send()
+        .await
+        .expect("should succeed");
+    match_events!(ev!(dns), ev!(connect), _, ev!(http(200)))(&mock.events());
+}
-- 
GitLab