LockLeaseService { Set<LockToken> refreshLockLeases(Set<LockToken> uncastedTokens) { if (uncastedTokens.isEmpty()) { return uncastedTokens; } LeaderTime leaderTime = time.get(); Set<LeasedLockToken> allTokens = leasedTokens(uncastedTokens); Set<LeasedLockToken> validByLease = allTokens.stream() .filter(token -> token.isValid(leaderTime)) .collect(Collectors.toSet()); Set<LeasedLockToken> toRefresh = Sets.difference(allTokens, validByLease); Set<LeasedLockToken> refreshedTokens = refreshTokens(toRefresh); return Sets.union(refreshedTokens, validByLease); } @VisibleForTesting  LockLeaseService(
            NamespacedConjureTimelockService delegate,
            UUID clientId);   }
@Test public void shouldOnlyCallIdentifiedTimeIfLeaseIsValid() { LeasedLockToken validToken = LeasedLockToken.of(LOCK_TOKEN, getLease()); when(timelock.leaderTime()).thenReturn(getIdentifiedTime()); lockLeaseService.refreshLockLeases(ImmutableSet.of(validToken)); verify(timelock).leaderTime(); verifyNoMoreInteractions(timelock); }
@Test public void shouldRefreshTheLease_invalidOnClient_validOnServer() { LeasedLockToken leasedLockToken = LeasedLockToken.of(LOCK_TOKEN, getLease(Duration.ZERO)); assertInvalid(leasedLockToken); when(timelock.refreshLocks(ConjureRefreshLocksRequest.of(ImmutableSet.of(LOCK_TOKEN)))) .thenReturn(ConjureRefreshLocksResponse.of( ImmutableSet.of(LOCK_TOKEN), getLease())); Set<LockToken> refreshed = lockLeaseService.refreshLockLeases(ImmutableSet.of(leasedLockToken)); verify(timelock).refreshLocks(ConjureRefreshLocksRequest.of(ImmutableSet.of(leasedLockToken.serverToken()))); LeasedLockToken refreshedLeasedLockToken = (LeasedLockToken) refreshed.iterator().next(); assertValid(refreshedLeasedLockToken); assertValid(leasedLockToken); assertThat(refreshedLeasedLockToken).isEqualTo(leasedLockToken); }