<commit_before>// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.container.jdisc.athenz.impl;

import com.yahoo.container.core.identity.IdentityConfig;
import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
import com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.RunnableWithTag;
import com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.Scheduler;
import com.yahoo.test.ManualClock;
import org.junit.Test;

import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.PriorityQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;

import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.INITIAL_BACKOFF_DELAY;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.INITIAL_WAIT_NTOKEN;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.MAX_REGISTER_BACKOFF_DELAY;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.REDUCED_UPDATE_PERIOD;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.REGISTER_INSTANCE_TAG;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.TIMEOUT_INITIAL_WAIT_TAG;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.UPDATE_CREDENTIALS_TAG;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.UPDATE_PERIOD;
import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

/**
 * @author mortent
 * @author bjorncs
 */
public class AthenzIdentityProviderImplTest {

    private static final IdentityConfig IDENTITY_CONFIG =
            new IdentityConfig(new IdentityConfig.Builder()
                                       .service("tenantService").domain("tenantDomain").loadBalancerAddress("cfg"));

    @Test
    public void athenz_credentials_are_retrieved_after_component_contruction_completed() {
        ServiceProviderApi serviceProviderApi = mock(ServiceProviderApi.class);
        AthenzService athenzService = mock(AthenzService.class);
        MockScheduler scheduler = new MockScheduler();

        when(serviceProviderApi.getSignedIdentityDocument()).thenReturn(getIdentityDocument());
        when(athenzService.sendInstanceRegisterRequest(any(), any())).thenReturn(
                new InstanceIdentity(null, "TOKEN"));
        AthenzCredentialsService credentialService =
                new AthenzCredentialsService(IDENTITY_CONFIG, serviceProviderApi, athenzService, scheduler.clock());

        AthenzIdentityProvider identityProvider =
                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, credentialService, scheduler, scheduler.clock());

        List<MockScheduler.CompletedTask> expectedTasks =
                Arrays.asList(
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, Duration.ZERO),
                        new MockScheduler.CompletedTask(TIMEOUT_INITIAL_WAIT_TAG, INITIAL_WAIT_NTOKEN));
        // Don't run update credential tasks, otherwise infinite loop
        List<MockScheduler.CompletedTask> completedTasks =
                scheduler.runAllTasks(task -> !task.tag().equals(UPDATE_CREDENTIALS_TAG));
        assertEquals(expectedTasks, completedTasks);
        assertEquals("TOKEN", identityProvider.getNToken());
    }

    @Test
    public void register_instance_uses_exponential_backoff() {
        AthenzCredentialsService credentialService = mock(AthenzCredentialsService.class);
        when(credentialService.registerInstance())
                .thenThrow(new RuntimeException("#1"))
                .thenThrow(new RuntimeException("#2"))
                .thenThrow(new RuntimeException("#3"))
                .thenThrow(new RuntimeException("#4"))
                .thenThrow(new RuntimeException("#5"))
                .thenReturn(new AthenzCredentials("TOKEN", null, null, null, Instant.now()));

        MockScheduler scheduler = new MockScheduler();
        AthenzIdentityProvider identityProvider =
                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, credentialService, scheduler, scheduler.clock());

        List<MockScheduler.CompletedTask> expectedTasks =
                Arrays.asList(
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, Duration.ZERO),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY),
                        new MockScheduler.CompletedTask(TIMEOUT_INITIAL_WAIT_TAG, INITIAL_WAIT_NTOKEN),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY.multipliedBy(2)),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY.multipliedBy(4)),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY.multipliedBy(8)),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, MAX_REGISTER_BACKOFF_DELAY));
        // Don't run update credential tasks, otherwise infinite loop
        List<MockScheduler.CompletedTask> completedTasks =
                scheduler.runAllTasks(task -> !task.tag().equals(UPDATE_CREDENTIALS_TAG));
        assertEquals(expectedTasks, completedTasks);
        assertEquals("TOKEN", identityProvider.getNToken());
    }

    @Test
    public void failed_credentials_updates_will_schedule_retries() {
        ServiceProviderApi serviceProviderApi = mock(ServiceProviderApi.class);
        AthenzService athenzService = mock(AthenzService.class);
        MockScheduler scheduler = new MockScheduler();

        when(serviceProviderApi.getSignedIdentityDocument()).thenReturn(getIdentityDocument());
        when(athenzService.sendInstanceRegisterRequest(any(), any())).thenReturn(
                new InstanceIdentity(null, "TOKEN"));
        when(athenzService.sendInstanceRefreshRequest(anyString(), anyString(), anyString(),
                                                      anyString(), any(), any(), any(), any()))
                .thenThrow(new RuntimeException("#1"))
                .thenThrow(new RuntimeException("#2"))
                .thenThrow(new RuntimeException("#3"))
                .thenReturn(new InstanceIdentity(null, "TOKEN"));
        AthenzCredentialsService credentialService =
                new AthenzCredentialsService(IDENTITY_CONFIG, serviceProviderApi, athenzService, scheduler.clock());

        AthenzIdentityProvider identityProvider =
                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, credentialService, scheduler, scheduler.clock());

        List<MockScheduler.CompletedTask> expectedTasks =
                Arrays.asList(
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, Duration.ZERO),
                        new MockScheduler.CompletedTask(TIMEOUT_INITIAL_WAIT_TAG, INITIAL_WAIT_NTOKEN),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, REDUCED_UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, REDUCED_UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, UPDATE_PERIOD));
        AtomicInteger counter = new AtomicInteger(0);
        List<MockScheduler.CompletedTask> completedTasks =
                scheduler.runAllTasks(task -> counter.getAndIncrement() < 7); // 1 registration + 1 timeout + 5 update tasks
        assertEquals(expectedTasks, completedTasks);
        assertEquals("TOKEN", identityProvider.getNToken());
    }

    private static String getIdentityDocument() {
        return "{\n" +
               "  \"identity-document\": \"<KEY>Q==\",\n" +
               "  \"signature\": \"kkEJB/98cy1FeXxzSjtvGH2a6BFgZu/9/kzCcAqRMZjENxnw5jyO1/bjZVzw2Sz4YHPsWSx2uxb32hiQ0U8rMP0zfA9nERIalSP0jB/hMU8laezGhdpk6VKZPJRC6YKAB9Bsv2qUIfMsSxkMqf66GUvjZAGaYsnNa2yHc1jIYHOGMeJO+HNPYJjGv26xPfAOPIKQzs3RmKrc3FoweTCsIwm5oblqekdJvVWYe0obwlOSB5uwc1zpq3Ie1QBFtJRuCGMVHg1pDPxXKBHLClGIrEvzLmICy6IRdHszSO5qiwujUD7sbrbM0sB/u0cYucxbcsGRUmBvme3UAw2mW9POVQ==\",\n" +
               "  \"signing-key-version\": 0,\n" +
               "  \"provider-unique-id\": \"tenant.application.dev.us-north-1.default.default.0\",\n" +
               "  \"dns-suffix\": \"dnsSuffix\",\n" +
               "  \"provider-service\": \"service\",\n" +
               "  \"zts-endpoint\": \"localhost/zts\", \n" +
               "  \"document-version\": 1\n" +
               "}";

    }

    private static class MockScheduler implements Scheduler {

        private final PriorityQueue<DelayedTask> tasks = new PriorityQueue<>();
        private final ManualClock clock = new ManualClock(Instant.EPOCH);

        @Override
        public void schedule(RunnableWithTag task, Duration delay) {
            tasks.offer(new DelayedTask(task, delay, clock.instant().plus(delay)));
        }

        List<CompletedTask> runAllTasks(Predicate<RunnableWithTag> filter) {
            List<CompletedTask> completedTasks = new ArrayList<>();
            while (!tasks.isEmpty()) {
                DelayedTask task = tasks.poll();
                RunnableWithTag runnable = task.runnableWithTag;
                if (filter.test(runnable)) {
                    clock.setInstant(task.startTime);
                    runnable.run();
                    completedTasks.add(new CompletedTask(runnable.tag(), task.delay));
                }
            }
            return completedTasks;
        }

        public ManualClock clock() {
            return clock;
        }

        private static class DelayedTask implements Comparable<DelayedTask> {
            final RunnableWithTag runnableWithTag;
            final Duration delay;
            final Instant startTime;<commit_msg>Take ManualClock as constructor argument<commit_after>// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.container.jdisc.athenz.impl;

import com.yahoo.container.core.identity.IdentityConfig;
import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
import com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.RunnableWithTag;
import com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.Scheduler;
import com.yahoo.test.ManualClock;
import org.junit.Test;

import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.PriorityQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;

import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.INITIAL_BACKOFF_DELAY;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.INITIAL_WAIT_NTOKEN;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.MAX_REGISTER_BACKOFF_DELAY;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.REDUCED_UPDATE_PERIOD;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.REGISTER_INSTANCE_TAG;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.TIMEOUT_INITIAL_WAIT_TAG;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.UPDATE_CREDENTIALS_TAG;
import static com.yahoo.container.jdisc.athenz.impl.AthenzIdentityProviderImpl.UPDATE_PERIOD;
import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

/**
 * @author mortent
 * @author bjorncs
 */
public class AthenzIdentityProviderImplTest {

    private static final IdentityConfig IDENTITY_CONFIG =
            new IdentityConfig(new IdentityConfig.Builder()
                                       .service("tenantService").domain("tenantDomain").loadBalancerAddress("cfg"));

    @Test
    public void athenz_credentials_are_retrieved_after_component_contruction_completed() {
        ServiceProviderApi serviceProviderApi = mock(ServiceProviderApi.class);
        AthenzService athenzService = mock(AthenzService.class);
        ManualClock clock = new ManualClock(Instant.EPOCH);
        MockScheduler scheduler = new MockScheduler(clock);

        when(serviceProviderApi.getSignedIdentityDocument()).thenReturn(getIdentityDocument());
        when(athenzService.sendInstanceRegisterRequest(any(), any())).thenReturn(
                new InstanceIdentity(null, "TOKEN"));
        AthenzCredentialsService credentialService =
                new AthenzCredentialsService(IDENTITY_CONFIG, serviceProviderApi, athenzService, clock);

        AthenzIdentityProvider identityProvider =
                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, credentialService, scheduler, clock);

        List<MockScheduler.CompletedTask> expectedTasks =
                Arrays.asList(
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, Duration.ZERO),
                        new MockScheduler.CompletedTask(TIMEOUT_INITIAL_WAIT_TAG, INITIAL_WAIT_NTOKEN));
        // Don't run update credential tasks, otherwise infinite loop
        List<MockScheduler.CompletedTask> completedTasks =
                scheduler.runAllTasks(task -> !task.tag().equals(UPDATE_CREDENTIALS_TAG));
        assertEquals(expectedTasks, completedTasks);
        assertEquals("TOKEN", identityProvider.getNToken());
    }

    @Test
    public void register_instance_uses_exponential_backoff() {
        AthenzCredentialsService credentialService = mock(AthenzCredentialsService.class);
        when(credentialService.registerInstance())
                .thenThrow(new RuntimeException("#1"))
                .thenThrow(new RuntimeException("#2"))
                .thenThrow(new RuntimeException("#3"))
                .thenThrow(new RuntimeException("#4"))
                .thenThrow(new RuntimeException("#5"))
                .thenReturn(new AthenzCredentials("TOKEN", null, null, null, Instant.now()));

        ManualClock clock = new ManualClock(Instant.EPOCH);
        MockScheduler scheduler = new MockScheduler(clock);
        AthenzIdentityProvider identityProvider =
                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, credentialService, scheduler, clock);

        List<MockScheduler.CompletedTask> expectedTasks =
                Arrays.asList(
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, Duration.ZERO),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY),
                        new MockScheduler.CompletedTask(TIMEOUT_INITIAL_WAIT_TAG, INITIAL_WAIT_NTOKEN),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY.multipliedBy(2)),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY.multipliedBy(4)),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, INITIAL_BACKOFF_DELAY.multipliedBy(8)),
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, MAX_REGISTER_BACKOFF_DELAY));
        // Don't run update credential tasks, otherwise infinite loop
        List<MockScheduler.CompletedTask> completedTasks =
                scheduler.runAllTasks(task -> !task.tag().equals(UPDATE_CREDENTIALS_TAG));
        assertEquals(expectedTasks, completedTasks);
        assertEquals("TOKEN", identityProvider.getNToken());
    }

    @Test
    public void failed_credentials_updates_will_schedule_retries() {
        ServiceProviderApi serviceProviderApi = mock(ServiceProviderApi.class);
        AthenzService athenzService = mock(AthenzService.class);
        ManualClock clock = new ManualClock(Instant.EPOCH);
        MockScheduler scheduler = new MockScheduler(clock);

        when(serviceProviderApi.getSignedIdentityDocument()).thenReturn(getIdentityDocument());
        when(athenzService.sendInstanceRegisterRequest(any(), any())).thenReturn(
                new InstanceIdentity(null, "TOKEN"));
        when(athenzService.sendInstanceRefreshRequest(anyString(), anyString(), anyString(),
                                                      anyString(), any(), any(), any(), any()))
                .thenThrow(new RuntimeException("#1"))
                .thenThrow(new RuntimeException("#2"))
                .thenThrow(new RuntimeException("#3"))
                .thenReturn(new InstanceIdentity(null, "TOKEN"));
        AthenzCredentialsService credentialService =
                new AthenzCredentialsService(IDENTITY_CONFIG, serviceProviderApi, athenzService, clock);

        AthenzIdentityProvider identityProvider =
                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, credentialService, scheduler, clock);

        List<MockScheduler.CompletedTask> expectedTasks =
                Arrays.asList(
                        new MockScheduler.CompletedTask(REGISTER_INSTANCE_TAG, Duration.ZERO),
                        new MockScheduler.CompletedTask(TIMEOUT_INITIAL_WAIT_TAG, INITIAL_WAIT_NTOKEN),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, REDUCED_UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, REDUCED_UPDATE_PERIOD),
                        new MockScheduler.CompletedTask(UPDATE_CREDENTIALS_TAG, UPDATE_PERIOD));
        AtomicInteger counter = new AtomicInteger(0);
        List<MockScheduler.CompletedTask> completedTasks =
                scheduler.runAllTasks(task -> counter.getAndIncrement() < 7); // 1 registration + 1 timeout + 5 update tasks
        assertEquals(expectedTasks, completedTasks);
        assertEquals("TOKEN", identityProvider.getNToken());
    }

    private static String getIdentityDocument() {
        return "{\n" +
               "  \"identity-document\": \"<KEY>" +
               "  \"signature\": \"<KEY>\",\n" +
               "  \"signing-key-version\": 0,\n" +
               "  \"provider-unique-id\": \"tenant.application.dev.us-north-1.default.default.0\",\n" +
               "  \"dns-suffix\": \"dnsSuffix\",\n" +
               "  \"provider-service\": \"service\",\n" +
               "  \"zts-endpoint\": \"localhost/zts\", \n" +
               "  \"document-version\": 1\n" +
               "}";

    }

    private static class MockScheduler implements Scheduler {

        private final PriorityQueue<DelayedTask> tasks = new PriorityQueue<>();
        private final ManualClock clock;

        MockScheduler(ManualClock clock) {
            this.clock = clock;
        }

        @Override
        public void schedule(RunnableWithTag task, Duration delay) {
            tasks.offer(new DelayedTask(task, delay, clock.instant().plus(delay)));
        }

        List<CompletedTask> runAllTasks(Predicate<RunnableWithTag> filter) {
            List<CompletedTask> completedTasks = new ArrayList<>();
            while (!tasks.isEmpty()) {
                DelayedTask task = tasks.poll();
                RunnableWithTag runnable = task.runnableWithTag;
                if (filter.test(runnable)) {
                    clock.setInstant(task.startTime);
                    runnable.run();
                    completedTasks.add(new CompletedTask(runnable.tag(), task.delay));
                }
            }
            return completedTasks;
        }

        private static class DelayedTask implements Comparable<DelayedTask> {
            final RunnableWithTag runnableWithTag;
            final Duration delay;
            final Instant startTime;
{'fraction_non_alphanumeric': 0.09225255302435192, 'fraction_numerical': 0.0048114689709348, 'mean_word_length': 3.0893394900622364, 'pattern_counts': {'":': 16, '<': 34, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4269682', 'n_tokens_mistral': 5743, 'n_tokens_neox': 5495, 'n_words': 876}