From e0eab8b0147e69f9a8e64a7878f129ae9e784d1f Mon Sep 17 00:00:00 2001
From: Jesse Szwedko <jesse@szwedko.me>
Date: Wed, 27 Apr 2022 07:32:36 -0700
Subject: [PATCH] Fix the rest of the usages

Signed-off-by: Jesse Szwedko <jesse@szwedko.me>
---
 src/aws/auth.rs                               | 42 +++++++++++++++----
 .../aws_kinesis_firehose/integration_tests.rs |  4 +-
 src/sinks/elasticsearch/integration_tests.rs  |  8 +++-
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/src/aws/auth.rs b/src/aws/auth.rs
index e4cfdb996c0b0..b2f727922b544 100644
--- a/src/aws/auth.rs
+++ b/src/aws/auth.rs
@@ -3,10 +3,11 @@ use aws_config::{
 };
 use aws_types::{credentials::SharedCredentialsProvider, region::Region, Credentials};
 use serde::{Deserialize, Serialize};
+use std::time::Duration;
 
 // matches default load timeout from the SDK as of 0.10.1, but lets us confidently document the
 // default rather than relying on the SDK default to not change
-const DEFAULT_LOAD_TIMEOUT_SECS: u64 = 5;
+const DEFAULT_LOAD_TIMEOUT: Duration = Duration::from_secs(5);
 
 /// Configuration for configuring authentication strategy for AWS.
 #[derive(Serialize, Deserialize, Clone, Debug, Derivative)]
@@ -80,9 +81,11 @@ async fn default_credentials_provider(
 ) -> SharedCredentialsProvider {
     let chain = DefaultCredentialsChain::builder()
         .region(region)
-        .load_timeout(std::time::Duration::from_secs(
-            load_timeout_secs.unwrap_or(DEFAULT_LOAD_TIMEOUT_SECS),
-        ));
+        .load_timeout(
+            load_timeout_secs
+                .map(Duration::from_secs)
+                .unwrap_or(DEFAULT_LOAD_TIMEOUT),
+        );
 
     SharedCredentialsProvider::new(chain.build().await)
 }
@@ -106,7 +109,24 @@ mod tests {
         )
         .unwrap();
 
-        assert!(matches!(config.auth, AwsAuthentication::Default {}));
+        assert!(matches!(config.auth, AwsAuthentication::Default { .. }));
+    }
+
+    #[test]
+    fn parsing_default_with_load_timeout() {
+        let config = toml::from_str::<ComponentConfig>(
+            r#"
+            load_timeout_secs = 10
+        "#,
+        )
+        .unwrap();
+
+        assert!(matches!(
+            config.auth,
+            AwsAuthentication::Default {
+                load_timeout_secs: Some(10)
+            }
+        ));
     }
 
     #[test]
@@ -118,7 +138,7 @@ mod tests {
         )
         .unwrap();
 
-        assert!(matches!(config.auth, AwsAuthentication::Default {}));
+        assert!(matches!(config.auth, AwsAuthentication::Default { .. }));
     }
 
     #[test]
@@ -126,6 +146,7 @@ mod tests {
         let config = toml::from_str::<ComponentConfig>(
             r#"
             auth.assume_role = "root"
+            auth.load_timeout_secs = 10
         "#,
         )
         .unwrap();
@@ -139,12 +160,19 @@ mod tests {
             r#"
             assume_role = "root"
             auth.assume_role = "auth.root"
+            auth.load_timeout_secs = 10
         "#,
         )
         .unwrap();
 
         match config.auth {
-            AwsAuthentication::Role { assume_role } => assert_eq!(&assume_role, "auth.root"),
+            AwsAuthentication::Role {
+                assume_role,
+                load_timeout_secs,
+            } => {
+                assert_eq!(&assume_role, "auth.root");
+                assert_eq!(load_timeout_secs, Some(10));
+            }
             _ => panic!(),
         }
     }
diff --git a/src/sinks/aws_kinesis_firehose/integration_tests.rs b/src/sinks/aws_kinesis_firehose/integration_tests.rs
index d5736af6aa4b0..15da5e62b3999 100644
--- a/src/sinks/aws_kinesis_firehose/integration_tests.rs
+++ b/src/sinks/aws_kinesis_firehose/integration_tests.rs
@@ -78,7 +78,9 @@ async fn firehose_put_records() {
     components::SINK_TESTS.assert(&AWS_SINK_TAGS);
 
     let config = ElasticsearchConfig {
-        auth: Some(ElasticsearchAuth::Aws(AwsAuthentication::Default {})),
+        auth: Some(ElasticsearchAuth::Aws(AwsAuthentication::Default {
+            load_timeout_secs: Some(5),
+        })),
         endpoint: elasticsearch_address(),
         bulk: Some(BulkConfig {
             index: Some(stream.clone()),
diff --git a/src/sinks/elasticsearch/integration_tests.rs b/src/sinks/elasticsearch/integration_tests.rs
index 721346a73caef..368acda4ad022 100644
--- a/src/sinks/elasticsearch/integration_tests.rs
+++ b/src/sinks/elasticsearch/integration_tests.rs
@@ -241,7 +241,9 @@ async fn insert_events_on_aws() {
 
     run_insert_tests(
         ElasticsearchConfig {
-            auth: Some(ElasticsearchAuth::Aws(AwsAuthentication::Default {})),
+            auth: Some(ElasticsearchAuth::Aws(AwsAuthentication::Default {
+                load_timeout_secs: Some(5),
+            })),
             endpoint: aws_server(),
             aws: Some(RegionOrEndpoint::with_region(String::from("localstack"))),
             ..config()
@@ -258,7 +260,9 @@ async fn insert_events_on_aws_with_compression() {
 
     run_insert_tests(
         ElasticsearchConfig {
-            auth: Some(ElasticsearchAuth::Aws(AwsAuthentication::Default {})),
+            auth: Some(ElasticsearchAuth::Aws(AwsAuthentication::Default {
+                load_timeout_secs: Some(5),
+            })),
             endpoint: aws_server(),
             aws: Some(RegionOrEndpoint::with_region(String::from("localstack"))),
             compression: Compression::gzip_default(),