Use Clock

2 files changed, 8 insertions, 9 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java
index 28e2fbadbce..93d2f9bf3b5 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/NTokenGenerator.java
@@ -6,10 +6,9 @@ import com.yahoo.athenz.auth.util.CryptoException;
 
 import java.security.PrivateKey;
 import java.security.SecureRandom;
+import java.time.Clock;
 import java.time.Duration;
-import java.time.Instant;
 import java.util.function.LongSupplier;
-import java.util.function.Supplier;
 
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
@@ -19,7 +18,7 @@ import static java.util.Objects.requireNonNull;
  */
 public class NTokenGenerator {
     private final Signer signer;
-    private final Supplier<Instant> time;
+    private final Clock clock;
 
     private String domain = null;
     private String name = null;
@@ -37,12 +36,12 @@ public class NTokenGenerator {
         String sign(String message, PrivateKey key) throws CryptoException;
     }
 
-    public NTokenGenerator() { this(Crypto::sign, Instant::now, new SecureRandom()::nextLong); }
+    public NTokenGenerator() { this(Crypto::sign, Clock.systemUTC(), new SecureRandom()::nextLong); }
 
     /** For testing. */
-    NTokenGenerator(Signer signer, Supplier<Instant> time, LongSupplier randomGenerator) {
+    NTokenGenerator(Signer signer, Clock clock, LongSupplier randomGenerator) {
         this.signer = signer;
-        this.time = time;
+        this.clock = clock;
         this.randomGenerator = randomGenerator;
     }
 
@@ -77,7 +76,7 @@ public class NTokenGenerator {
     public NToken sign(PrivateKey privateKey) {
         // See https://github.com/AthenZ/athenz/blob/master/libs/go/zmssvctoken/token.go
 
-        var generationTime = time.get();
+        var generationTime = clock.instant();
 
         token.setLength(0);
         append('v', "S1");
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java
index 0e07b90f4d6..7c4ee36926f 100644
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/api/NTokenGeneratorTest.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.athenz.api;
 
 import com.yahoo.athenz.auth.util.CryptoException;
+import com.yahoo.test.ManualClock;
 import org.junit.Test;
 
 import java.security.PrivateKey;
@@ -15,13 +16,12 @@ import static org.mockito.Mockito.mock;
  * @author hakonhall
  */
 public class NTokenGeneratorTest {
-    private NTokenGenerator generator;
     private final PrivateKey key = mock(PrivateKey.class);
 
     @Test
     public void ntoken() {
         var signer = new Signer("signature");
-        generator = new NTokenGenerator(signer, () -> Instant.ofEpochSecond(12L), () -> 3L);
+        NTokenGenerator generator = new NTokenGenerator(signer, new ManualClock(Instant.ofEpochSecond(12L)), () -> 3L);
         AthenzIdentity identity = new AthenzService("domain", "service");
 
         NToken token = generator.setIdentity(identity)