google.registry.testing.TaskQueueHelper.java Source code

Java tutorial

Introduction

Here is the source code for google.registry.testing.TaskQueueHelper.java

Source

// Copyright 2016 The Nomulus Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package google.registry.testing;

import static com.google.appengine.tools.development.testing.LocalTaskQueueTestConfig.getLocalTaskQueue;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Predicates.in;
import static com.google.common.base.Predicates.not;
import static com.google.common.collect.Iterables.getFirst;
import static com.google.common.collect.Iterables.transform;
import static com.google.common.collect.Multisets.containsOccurrences;
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assert_;
import static google.registry.util.DiffUtils.prettyPrintEntityDeepDiff;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Arrays.asList;

import com.google.appengine.api.taskqueue.dev.QueueStateInfo;
import com.google.appengine.api.taskqueue.dev.QueueStateInfo.HeaderWrapper;
import com.google.appengine.api.taskqueue.dev.QueueStateInfo.TaskStateInfo;
import com.google.common.base.Ascii;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.Predicate;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.ImmutableMultiset;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.net.HttpHeaders;
import com.google.common.net.MediaType;
import google.registry.dns.DnsConstants;
import google.registry.model.ImmutableObject;
import java.net.URI;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import javax.annotation.Nonnull;
import org.joda.time.Duration;

/** Static utility functions for testing task queues. */
public class TaskQueueHelper {

    /**
     * Matcher to match against the tasks in the task queue. Fields that aren't set are not compared.
     */
    public static class TaskMatcher implements Predicate<TaskStateInfo> {

        MatchableTaskInfo expected = new MatchableTaskInfo();

        public TaskMatcher taskName(String taskName) {
            expected.taskName = taskName;
            return this;
        }

        public TaskMatcher url(String url) {
            expected.url = url;
            return this;
        }

        public TaskMatcher method(String method) {
            expected.method = method;
            return this;
        }

        public TaskMatcher payload(String payload) {
            checkState(expected.params.isEmpty(), "Cannot add a payload to a TaskMatcher with params.");
            expected.payload = payload;
            return this;
        }

        public TaskMatcher tag(String tag) {
            expected.tag = tag;
            return this;
        }

        public TaskMatcher header(String name, String value) {
            // Lowercase for case-insensitive comparison.
            expected.headers.put(Ascii.toLowerCase(name), value);
            return this;
        }

        public TaskMatcher param(String key, String value) {
            checkState(expected.payload == null, "Cannot add params to a TaskMatcher with a payload.");
            expected.params.put(key, value);
            return this;
        }

        public TaskMatcher etaDelta(Duration lowerBound, Duration upperBound) {
            checkState(!lowerBound.isShorterThan(Duration.ZERO), "lowerBound must be non-negative.");
            checkState(upperBound.isLongerThan(lowerBound), "upperBound must be greater than lowerBound.");
            expected.etaDeltaLowerBound = lowerBound.getStandardSeconds();
            expected.etaDeltaUpperBound = upperBound.getStandardSeconds();
            return this;
        }

        /**
         * Returns {@code true} if there are not more occurrences in {@code sub} of each of its entries
         * than there are in {@code super}.
         */
        private static boolean containsEntries(Multimap<?, ?> superMultimap, Multimap<?, ?> subMultimap) {
            return containsOccurrences(ImmutableMultiset.copyOf(superMultimap.entries()),
                    ImmutableMultiset.copyOf(subMultimap.entries()));
        }

        /**
         * Returns true if the fields set on the current object match the given task info. This is not
         * quite the same contract as {@link #equals}, since it will ignore null fields.
         *
         * <p>Match fails if any headers or params expected on the TaskMatcher are not found on the
         * TaskStateInfo. Note that the inverse is not true (i.e. there may be extra headers on the
         * TaskStateInfo).
         */
        @Override
        public boolean apply(@Nonnull TaskStateInfo info) {
            MatchableTaskInfo actual = new MatchableTaskInfo(info);
            return (expected.taskName == null || Objects.equals(expected.taskName, actual.taskName))
                    && (expected.url == null || Objects.equals(expected.url, actual.url))
                    && (expected.method == null || Objects.equals(expected.method, actual.method))
                    && (expected.payload == null || Objects.equals(expected.payload, actual.payload))
                    && (expected.tag == null || Objects.equals(expected.tag, actual.tag))
                    && (expected.etaDeltaLowerBound == null || expected.etaDeltaLowerBound <= actual.etaDelta)
                    && (expected.etaDeltaUpperBound == null || expected.etaDeltaUpperBound >= actual.etaDelta)
                    && containsEntries(actual.params, expected.params)
                    && containsEntries(actual.headers, expected.headers);
        }

        @Override
        public String toString() {
            return Joiner.on('\n').withKeyValueSeparator(":\n")
                    .join(Maps.transformValues(expected.toMap(), new Function<Object, String>() {
                        @Override
                        public String apply(Object input) {
                            return "\t" + String.valueOf(input).replaceAll("\n", "\n\t");
                        }
                    }));
        }
    }

    /** Returns the info object for the provided queue name. */
    public static QueueStateInfo getQueueInfo(String queueName) {
        return getLocalTaskQueue().getQueueStateInfo().get(queueName);
    }

    /**
     * Ensures that the tasks in the named queue are exactly those with the expected property
     * values after being transformed with the provided property getter function.
     */
    public static void assertTasksEnqueuedWithProperty(String queueName,
            Function<TaskStateInfo, String> propertyGetter, String... expectedTaskProperties) throws Exception {
        // Ordering is irrelevant but duplicates should be considered independently.
        assertThat(transform(getQueueInfo(queueName).getTaskInfo(), propertyGetter))
                .containsExactly((Object[]) expectedTaskProperties);
    }

    /** Ensures that the tasks in the named queue are exactly those with the expected names. */
    public static void assertTasksEnqueued(String queueName, String... expectedTaskNames) throws Exception {
        Function<TaskStateInfo, String> nameGetter = new Function<TaskStateInfo, String>() {
            @Nonnull
            @Override
            public String apply(@Nonnull TaskStateInfo taskStateInfo) {
                return taskStateInfo.getTaskName();
            }
        };
        assertTasksEnqueuedWithProperty(queueName, nameGetter, expectedTaskNames);
    }

    /**
     * Ensures that the only tasks in the named queue are exactly those that match the expected
     * matchers.
     */
    public static void assertTasksEnqueued(String queueName, TaskMatcher... taskMatchers) throws Exception {
        assertTasksEnqueued(queueName, Arrays.asList(taskMatchers));
    }

    /**
     * Ensures that the only tasks in the named queue are exactly those that match the expected
     * matchers.
     */
    public static void assertTasksEnqueued(String queueName, List<TaskMatcher> taskMatchers) throws Exception {
        QueueStateInfo qsi = getQueueInfo(queueName);
        assertThat(qsi.getTaskInfo()).hasSize(taskMatchers.size());
        LinkedList<TaskStateInfo> taskInfos = new LinkedList<>(qsi.getTaskInfo());
        for (final TaskMatcher taskMatcher : taskMatchers) {
            try {
                taskInfos.remove(Iterables.find(taskInfos, taskMatcher));
            } catch (NoSuchElementException e) {
                final Map<String, Object> taskMatcherMap = taskMatcher.expected.toMap();
                assert_().fail("Task not found in queue %s:\n\n%s\n\nPotential candidate match diffs:\n\n%s",
                        queueName, taskMatcher,
                        FluentIterable.from(taskInfos).transform(new Function<TaskStateInfo, String>() {
                            @Override
                            public String apply(TaskStateInfo input) {
                                return prettyPrintEntityDeepDiff(taskMatcherMap, Maps.filterKeys(
                                        new MatchableTaskInfo(input).toMap(), in(taskMatcherMap.keySet())));
                            }
                        }).join(Joiner.on('\n')));
            }
        }
    }

    /** Empties the task queue. */
    public static void clearTaskQueue(String queueName) throws Exception {
        getLocalTaskQueue().flushQueue(queueName);
    }

    /** Asserts at least one task exists in {@code queue}. */
    public static void assertAtLeastOneTaskIsEnqueued(String queue) throws Exception {
        assertThat(getQueueInfo(queue).getCountTasks()).isGreaterThan(0);
    }

    /** Ensures that the named queue contains no tasks. */
    public static void assertNoTasksEnqueued(String... queueNames) throws Exception {
        for (String queueName : queueNames) {
            assertThat(getQueueInfo(queueName).getCountTasks()).isEqualTo(0);
        }
    }

    /** Returns the value for the param on a task info, or empty if it is missing. */
    private static String getParamFromTaskInfo(TaskStateInfo taskInfo, String paramName) {
        return getFirst(UriParameters.parse(taskInfo.getBody()).get(paramName), "");
    }

    /** Ensures that the DNS queue tasks are exactly those for the expected target names. */
    public static void assertDnsTasksEnqueued(String... expectedTaskTargetNames) throws Exception {
        assertTasksEnqueuedWithProperty(DnsConstants.DNS_PULL_QUEUE_NAME, new Function<TaskStateInfo, String>() {
            @Nonnull
            @Override
            public String apply(@Nonnull TaskStateInfo taskStateInfo) {
                return getParamFromTaskInfo(taskStateInfo, DnsConstants.DNS_TARGET_NAME_PARAM);
            }
        }, expectedTaskTargetNames);
    }

    /** Ensures that the DNS queue does not contain any tasks. */
    public static void assertNoDnsTasksEnqueued() throws Exception {
        assertNoTasksEnqueued(DnsConstants.DNS_PULL_QUEUE_NAME);
    }

    /** An adapter to clean up a {@link TaskStateInfo} for ease of matching. */
    private static class MatchableTaskInfo extends ImmutableObject {

        String taskName;
        String method;
        String url;
        String payload;
        String tag;
        Double etaDelta;
        Long etaDeltaLowerBound;
        Long etaDeltaUpperBound;
        Multimap<String, String> headers = ArrayListMultimap.create();
        Multimap<String, String> params = ArrayListMultimap.create();

        MatchableTaskInfo() {
        }

        MatchableTaskInfo(TaskStateInfo info) {
            URI uri;
            try {
                uri = new URI(info.getUrl());
            } catch (java.net.URISyntaxException e) {
                throw new IllegalArgumentException(e);
            }
            this.taskName = info.getTaskName();
            this.method = info.getMethod();
            this.url = uri.getPath();
            this.payload = info.getBody();
            this.etaDelta = info.getEtaDelta();
            if (info.getTagAsBytes() != null) {
                this.tag = new String(info.getTagAsBytes(), UTF_8);
            }
            ImmutableMultimap.Builder<String, String> headerBuilder = new ImmutableMultimap.Builder<>();
            for (HeaderWrapper header : info.getHeaders()) {
                // Lowercase header name for comparison since HTTP
                // header names are case-insensitive.
                headerBuilder.put(Ascii.toLowerCase(header.getKey()), header.getValue());
            }
            this.headers = headerBuilder.build();
            ImmutableMultimap.Builder<String, String> inputParams = new ImmutableMultimap.Builder<>();
            String query = uri.getQuery();
            if (query != null) {
                inputParams.putAll(UriParameters.parse(query));
            }
            if (headers.containsEntry(Ascii.toLowerCase(HttpHeaders.CONTENT_TYPE),
                    MediaType.FORM_DATA.toString())) {
                inputParams.putAll(UriParameters.parse(info.getBody()));
            }
            this.params = inputParams.build();
        }

        public Map<String, Object> toMap() {
            Map<String, Object> builder = new HashMap<>();
            builder.put("taskName", taskName);
            builder.put("url", url);
            builder.put("method", method);
            builder.put("headers", headers.asMap());
            builder.put("params", params.asMap());
            builder.put("payload", payload);
            builder.put("tag", tag);
            builder.put("etaDelta", etaDelta);
            builder.put("etaDeltaLowerBound", etaDeltaLowerBound);
            builder.put("etaDeltaUpperBound", etaDeltaUpperBound);
            return Maps.filterValues(builder, not(in(asList(null, "", Collections.EMPTY_MAP))));
        }
    }
}