Fix refresh logic

2 files changed, 39 insertions, 7 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialsProvider.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialsProvider.java
index 28f028832b4..bd2f76bac52 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialsProvider.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialsProvider.java
@@ -15,8 +15,7 @@ import javax.net.ssl.SSLContext;
 import java.net.URI;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.Objects;
-import java.util.logging.Logger;
+import java.util.Optional;
 
 /**
  * Implementation of AWSCredentialsProvider using com.yahoo.vespa.athenz.client.zts.ZtsClient
@@ -25,8 +24,6 @@ import java.util.logging.Logger;
  */
 public class AwsCredentialsProvider implements AWSCredentialsProvider {
 
-    private static final Logger logger = Logger.getLogger(AwsCredentialsProvider.class.getName());
-
     private final static Duration MIN_EXPIRY = Duration.ofMinutes(5);
     private final AthenzDomain athenzDomain;
     private final AwsRole awsRole;
@@ -72,8 +69,8 @@ public class AwsCredentialsProvider implements AWSCredentialsProvider {
     /*
      * Checks credential expiration, returns true if it will expipre in the next MIN_EXPIRY minutes
      */
-    private static boolean shouldRefresh(AwsTemporaryCredentials credentials) {
-        Instant expiration = credentials.expiration();
-        return Objects.isNull(expiration) || expiration.minus(MIN_EXPIRY).isAfter(Instant.now());
+    static boolean shouldRefresh(AwsTemporaryCredentials credentials) {
+        Instant expiration = Optional.ofNullable(credentials).map(AwsTemporaryCredentials::expiration).orElse(Instant.EPOCH);
+        return Duration.between(Instant.now(), expiration).toMinutes() < MIN_EXPIRY.toMinutes();
     }
 }
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialProviderTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialProviderTest.java
new file mode 100644
index 00000000000..d637dcae14c
--- /dev/null
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/client/aws/AwsCredentialProviderTest.java
@@ -0,0 +1,35 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.athenz.client.aws;
+
+import com.yahoo.vespa.athenz.api.AwsTemporaryCredentials;
+import org.junit.Test;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class AwsCredentialProviderTest {
+
+    @Test
+    public void refreshes_correctly() {
+        Clock clock = Clock.systemUTC();
+        // Does not require refresh when expires in 10 minutes
+        assertFalse(AwsCredentialsProvider.shouldRefresh(getCredentials(clock.instant().plus(Duration.ofMinutes(10)))));
+
+        // Requires refresh when expires in 3 minutes
+        assertTrue(AwsCredentialsProvider.shouldRefresh(getCredentials(clock.instant().plus(Duration.ofMinutes(3)))));
+
+        // Requires refresh when expired
+        assertTrue(AwsCredentialsProvider.shouldRefresh(getCredentials(clock.instant().minus(Duration.ofMinutes(1)))));
+
+        // Refreshes when no credentials provided
+        assertTrue(AwsCredentialsProvider.shouldRefresh(null));
+    }
+
+    private AwsTemporaryCredentials getCredentials(Instant expiration) {
+        return new AwsTemporaryCredentials("accesskey", "secretaccesskey", "sessionToken", expiration);
+    }
+}