1 1 | /*
|
2 2 | * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
3 3 | * SPDX-License-Identifier: Apache-2.0
|
4 4 | */
|
5 5 |
|
6 6 | use aws_config::Region;
|
7 7 | use aws_credential_types::Credentials;
|
8 + | use aws_runtime::user_agent::test_util::assert_ua_contains_metric_values;
|
8 9 | use aws_sdk_dynamodb::{
|
9 10 | config::Builder,
|
10 11 | error::{DisplayErrorContext, SdkError},
|
11 12 | operation::batch_get_item::{BatchGetItemError, BatchGetItemOutput},
|
12 13 | types::{AttributeValue, KeysAndAttributes},
|
13 14 | Client, Config,
|
14 15 | };
|
15 16 | use aws_smithy_http_client::test_util::{capture_request, CaptureRequestReceiver};
|
16 17 | use aws_smithy_runtime::assert_str_contains;
|
17 18 | use aws_smithy_runtime_api::http::Response;
|
18 19 | use aws_types::endpoint_config::AccountIdEndpointMode;
|
19 20 |
|
20 21 | fn test_client(update_builder: fn(Builder) -> Builder) -> (Client, CaptureRequestReceiver) {
|
21 22 | let (http_client, request) = capture_request(None);
|
22 23 | let builder = Config::builder()
|
23 24 | .region(Region::new("us-east-1"))
|
24 25 | .credentials_provider(
|
25 26 | Credentials::builder()
|
26 27 | .account_id("123456789012")
|
27 28 | .access_key_id("ANOTREAL")
|
28 29 | .secret_access_key("notrealrnrELgWzOk3IfjzDKtFBhDby")
|
29 30 | .provider_name("test")
|
30 31 | .build(),
|
31 32 | )
|
32 33 | .http_client(http_client);
|
33 34 | (Client::from_conf(update_builder(builder).build()), request)
|
34 35 | }
|
35 36 |
|
36 37 | async fn call_operation(
|
37 38 | client: Client,
|
38 39 | table_name: &str,
|
39 40 | ) -> Result<BatchGetItemOutput, SdkError<BatchGetItemError, Response>> {
|
40 41 | let mut attr_v = std::collections::HashMap::new();
|
41 42 | attr_v.insert(":s".to_string(), AttributeValue::S("value".into()));
|
42 43 | let mut kv = std::collections::HashMap::new();
|
43 44 | kv.insert(":pk".to_string(), AttributeValue::M(attr_v));
|
44 45 | client
|
45 46 | .batch_get_item()
|
46 47 | .request_items(
|
47 48 | table_name,
|
48 49 | KeysAndAttributes::builder().keys(kv).build().unwrap(),
|
49 50 | )
|
50 51 | .send()
|
51 52 | .await
|
52 53 | }
|
53 54 |
|
54 55 | #[tokio::test]
|
55 56 | async fn basic_positive_cases() {
|
56 - | let test_cases: &[(fn(Builder) -> Builder, &str, &str)] = &[
|
57 + | let test_cases: &[(fn(Builder) -> Builder, &str, &str, &[&'static str])] = &[
|
57 58 | (
|
58 59 | std::convert::identity,
|
59 60 | "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
|
60 61 | "https://333333333333.ddb.us-east-1.amazonaws.com/",
|
62 + | &["P", "T"],
|
61 63 | ),
|
62 64 | (
|
63 65 | std::convert::identity,
|
64 66 | "table_name", // doesn't specify ARN for the table name
|
65 67 | "https://123456789012.ddb.us-east-1.amazonaws.com/", // the account ID should come from credentials
|
68 + | &["P", "T"],
|
66 69 | ),
|
67 70 | (
|
68 71 | |b: Builder| b.credentials_provider(Credentials::for_tests()), // credentials do not provide an account ID
|
69 72 | "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
|
70 73 | "https://333333333333.ddb.us-east-1.amazonaws.com/",
|
74 + | &["P"],
|
71 75 | ),
|
72 76 | (
|
73 77 | |b: Builder| b.account_id_endpoint_mode(AccountIdEndpointMode::Preferred), // sets the default mode `Preferred` explicitly
|
74 78 | "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
|
75 79 | "https://333333333333.ddb.us-east-1.amazonaws.com/",
|
80 + | &["P", "T"],
|
76 81 | ),
|
77 82 | (
|
78 83 | |b: Builder| b.account_id_endpoint_mode(AccountIdEndpointMode::Disabled),
|
79 84 | "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
|
80 85 | "https://dynamodb.us-east-1.amazonaws.com/",
|
86 + | &["Q", "T"],
|
81 87 | ),
|
82 88 | (
|
83 89 | |b: Builder| b.account_id_endpoint_mode(AccountIdEndpointMode::Required),
|
84 90 | "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
|
85 91 | "https://333333333333.ddb.us-east-1.amazonaws.com/",
|
92 + | &["R", "T"],
|
86 93 | ),
|
87 94 | ];
|
88 95 |
|
89 - | for (i, (update_builder, table_name, expected_uri)) in test_cases.into_iter().enumerate() {
|
96 + | for (i, (update_builder, table_name, expected_uri, expected_metrics)) in
|
97 + | test_cases.into_iter().enumerate()
|
98 + | {
|
90 99 | let (client, rx) = test_client(*update_builder);
|
91 100 | let _ = call_operation(client, table_name).await;
|
92 101 | let req = rx.expect_request();
|
93 102 | assert_eq!(
|
94 103 | *expected_uri,
|
95 104 | req.uri(),
|
96 105 | "on the {i}th test case where table name is `{table_name}`"
|
97 106 | );
|
107 + |
|
108 + | // Test User-Agent metrics for account ID tracking
|
109 + | let user_agent = req.headers().get("x-amz-user-agent").unwrap();
|
110 + | assert_ua_contains_metric_values(user_agent, expected_metrics);
|
98 111 | }
|
99 112 | }
|
100 113 |
|
101 114 | #[tokio::test]
|
102 115 | async fn error_should_be_raised_when_account_id_is_expected_but_not_resolved() {
|
103 116 | let (client, _) = test_client(|b| {
|
104 117 | b.account_id_endpoint_mode(AccountIdEndpointMode::Required)
|
105 118 | .credentials_provider(Credentials::for_tests())
|
106 119 | });
|
107 120 | // doesn't specify ARN for the table name
|