/* -*- Mode: Java; c-basic-offset: 4; tab-width: 4; indent-tabs-mode: nil; -*-
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package org.mozilla.geckoview.test.rule;

import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONTokener;
import org.mozilla.gecko.util.ThreadUtils;
import org.mozilla.geckoview.Autofill;
import org.mozilla.geckoview.ContentBlocking;
import org.mozilla.geckoview.GeckoDisplay;
import org.mozilla.geckoview.GeckoResult;
import org.mozilla.geckoview.GeckoRuntime;
import org.mozilla.geckoview.GeckoSession;
import org.mozilla.geckoview.GeckoSessionSettings;
import org.mozilla.geckoview.MediaSession;
import org.mozilla.geckoview.RuntimeTelemetry;
import org.mozilla.geckoview.SessionTextInput;
import org.mozilla.geckoview.WebExtension;
import org.mozilla.geckoview.WebExtensionController;
import org.mozilla.geckoview.test.util.TestServer;
import org.mozilla.geckoview.test.util.RuntimeCreator;
import org.mozilla.geckoview.test.util.Environment;
import org.mozilla.geckoview.test.util.UiThreadUtils;
import org.mozilla.geckoview.test.util.Callbacks;

import static org.hamcrest.Matchers.*;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;

import org.hamcrest.Matcher;

import org.json.JSONObject;

import org.junit.rules.ErrorCollector;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

import android.app.Instrumentation;
import android.graphics.Point;
import android.graphics.SurfaceTexture;
import android.os.SystemClock;
import androidx.annotation.NonNull;
import androidx.annotation.Nullable;
import androidx.test.platform.app.InstrumentationRegistry;
import android.util.Log;
import android.util.Pair;
import android.view.MotionEvent;
import android.view.Surface;

import java.io.File;
import java.lang.annotation.Annotation;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Proxy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;

import kotlin.jvm.JvmClassMappingKt;
import kotlin.reflect.KClass;

/**
 * TestRule that, for each test, sets up a GeckoSession, runs the test on the UI thread,
 * and tears down the GeckoSession at the end of the test. The rule also provides methods
 * for waiting on particular callbacks to be called, and methods for asserting that
 * callbacks are called in the proper order.
 */
public class GeckoSessionTestRule implements TestRule {
    private static final String LOGTAG = "GeckoSessionTestRule";

    private static final int TEST_PORT = 4245;
    public static final String TEST_ENDPOINT = "http://localhost:" + TEST_PORT;

    private static final Method sOnPageStart;
    private static final Method sOnPageStop;
    private static final Method sOnNewSession;
    private static final Method sOnCrash;
    private static final Method sOnKill;

    static {
        try {
            sOnPageStart = GeckoSession.ProgressDelegate.class.getMethod(
                    "onPageStart", GeckoSession.class, String.class);
            sOnPageStop = GeckoSession.ProgressDelegate.class.getMethod(
                    "onPageStop", GeckoSession.class, boolean.class);
            sOnNewSession = GeckoSession.NavigationDelegate.class.getMethod(
                    "onNewSession", GeckoSession.class, String.class);
            sOnCrash = GeckoSession.ContentDelegate.class.getMethod(
                    "onCrash", GeckoSession.class);
            sOnKill = GeckoSession.ContentDelegate.class.getMethod(
                    "onKill", GeckoSession.class);
        } catch (final NoSuchMethodException e) {
            throw new RuntimeException(e);
        }
    }

    public void addDisplay(final GeckoSession session, final int x, final int y) {
        final GeckoDisplay display = session.acquireDisplay();

        final SurfaceTexture displayTexture = new SurfaceTexture(0);
        displayTexture.setDefaultBufferSize(x, y);

        final Surface displaySurface = new Surface(displayTexture);
        display.surfaceChanged(displaySurface, x, y);

        mDisplays.put(session, display);
        mDisplayTextures.put(session, displayTexture);
        mDisplaySurfaces.put(session, displaySurface);
    }

    public void releaseDisplay(final GeckoSession session) {
        if (!mDisplays.containsKey(session)) {
            // No display to release
            return;
        }
        final GeckoDisplay display = mDisplays.remove(session);
        display.surfaceDestroyed();
        session.releaseDisplay(display);
        final Surface displaySurface = mDisplaySurfaces.remove(session);
        displaySurface.release();
        final SurfaceTexture displayTexture = mDisplayTextures.remove(session);
        displayTexture.release();
    }

    /**
     * Specify the timeout for any of the wait methods, in milliseconds, relative to
     * {@link Environment#DEFAULT_TIMEOUT_MILLIS}. When the default timeout scales to account
     * for differences in the device under test, the timeout value here will be
     * scaled as well. Can be used on classes or methods.
     */
    @Target({ElementType.METHOD, ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface TimeoutMillis {
        long value();
    }

    /**
     * Specify the display size for the GeckoSession in device pixels
     */
    @Target({ElementType.METHOD, ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface WithDisplay {
        int width();
        int height();
    }

    /**
     * Specify that the main session should not be opened at the start of the test.
     */
    @Target({ElementType.METHOD, ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface ClosedSessionAtStart {
        boolean value() default true;
    }

    /**
     * Specify that the test will set a delegate to null when creating a session, rather
     * than setting the delegate to a proxy. The test cannot wait on any delegates that
     * are set to null.
     */
    @Target({ElementType.METHOD, ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface NullDelegate {
        Class<?> value();

        @Target({ElementType.METHOD, ElementType.TYPE})
        @Retention(RetentionPolicy.RUNTIME)
        @interface List {
            NullDelegate[] value();
        }
    }

    /**
     * Specify a list of GeckoSession settings to be applied to the GeckoSession object
     * under test. Can be used on classes or methods. Note that the settings values must
     * be string literals regardless of the type of the settings.
     * <p>
     * Enable tracking protection for a particular test:
     * <pre>
     * &#64;Setting.List(&#64;Setting(key = Setting.Key.USE_TRACKING_PROTECTION,
     *                        value = "false"))
     * &#64;Test public void test() { ... }
     * </pre>
     * <p>
     * Use multiple settings:
     * <pre>
     * &#64;Setting.List({&#64;Setting(key = Setting.Key.USE_PRIVATE_MODE,
     *                         value = "true"),
     *                &#64;Setting(key = Setting.Key.USE_TRACKING_PROTECTION,
     *                         value = "false")})
     * </pre>
     */
    @Target({ElementType.METHOD, ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface Setting {
        enum Key {
            CHROME_URI,
            DISPLAY_MODE,
            ALLOW_JAVASCRIPT,
            SCREEN_ID,
            USE_PRIVATE_MODE,
            USE_TRACKING_PROTECTION,
            FULL_ACCESSIBILITY_TREE;

            private final GeckoSessionSettings.Key<?> mKey;
            private final Class<?> mType;

            Key() {
                final Field field;
                try {
                    field = GeckoSessionSettings.class.getDeclaredField(name());
                    field.setAccessible(true);
                    mKey = (GeckoSessionSettings.Key<?>) field.get(null);
                } catch (final NoSuchFieldException | IllegalAccessException e) {
                    throw new RuntimeException(e);
                }

                final ParameterizedType genericType = (ParameterizedType) field.getGenericType();
                mType = (Class<?>) genericType.getActualTypeArguments()[0];
            }

            @SuppressWarnings("unchecked")
            public void set(final GeckoSessionSettings settings, final String value) {
                try {
                    if (boolean.class.equals(mType) || Boolean.class.equals(mType)) {
                        final Method method = GeckoSessionSettings.class
                                .getDeclaredMethod("setBoolean",
                                        GeckoSessionSettings.Key.class,
                                        boolean.class);
                        method.setAccessible(true);
                        method.invoke(settings, mKey, Boolean.valueOf(value));
                    } else if (int.class.equals(mType) || Integer.class.equals(mType)) {
                        final Method method = GeckoSessionSettings.class
                                .getDeclaredMethod("setInt",
                                        GeckoSessionSettings.Key.class,
                                        int.class);
                        method.setAccessible(true);
                        try {
                            method.invoke(settings, mKey,
                                    (Integer)GeckoSessionSettings.class.getField(value)
                                            .get(null));
                        }
                        catch (final NoSuchFieldException | IllegalAccessException |
                                ClassCastException e) {
                            method.invoke(settings, mKey,
                                    Integer.valueOf(value));
                        }
                    } else if (String.class.equals(mType)) {
                        final Method method = GeckoSessionSettings.class
                                .getDeclaredMethod("setString",
                                        GeckoSessionSettings.Key.class,
                                        String.class);
                        method.setAccessible(true);
                        method.invoke(settings, mKey, value);
                    } else {
                        throw new IllegalArgumentException("Unsupported type: " +
                                mType.getSimpleName());
                    }
                } catch (final NoSuchMethodException
                        | IllegalAccessException
                        | InvocationTargetException e) {
                    throw new RuntimeException(e);
                }
            }
        }

        @Target({ElementType.METHOD, ElementType.TYPE})
        @Retention(RetentionPolicy.RUNTIME)
        @interface List {
            Setting[] value();
        }

        Key key();
        String value();
    }

    /**
     * Assert that a method is called or not called, and if called, the order and number
     * of times it is called. The order number is a monotonically increasing integer; if
     * an called method's order number is less than the current order number, an exception
     * is raised for out-of-order call.
     * <p>
     * {@code @AssertCalled} asserts the method must be called at least once.
     * <p>
     * {@code @AssertCalled(false)} asserts the method must not be called.
     * <p>
     * {@code @AssertCalled(order = 2)} asserts the method must be called once and
     * after any other method with order number less than 2.
     * <p>
     * {@code @AssertCalled(order = {2, 4})} asserts order number 2 for first
     * call and order number 4 for any subsequent calls.
     * <p>
     * {@code @AssertCalled(count = 2)} asserts two calls total in any order
     * with respect to other calls.
     * <p>
     * {@code @AssertCalled(count = 2, order = 2)} asserts two calls, both with
     * order number 2.
     * <p>
     * {@code @AssertCalled(count = 2, order = {2, 4, 6})} asserts two calls
     * total: the first with order number 2 and the second with order number 4.
     */
    @Target(ElementType.METHOD)
    @Retention(RetentionPolicy.RUNTIME)
    public @interface AssertCalled {
        /**
         * @return True if the method must be called if count != 0,
         *         or false if the method must not be called.
         */
        boolean value() default true;

        /**
         * @return The number of calls allowed. Specify -1 to allow any number > 0. Specify 0 to
         *         assert the method is not called, even if value() is true.
         */
        int count() default -1;

        /**
         * @return If called, the order number for each call, or 0 to allow arbitrary
         *         order. If order's length is more than count, extra elements are not used;
         *         if order's length is less than count, the last element is repeated.
         */
        int[] order() default 0;
    }

    /**
     * Interface that represents a function that registers or unregisters a delegate.
     */
    public interface DelegateRegistrar<T> {
        void invoke(T delegate) throws Throwable;
    }

    /*
     * If the value here is true, content crashes will be ignored. If false, the test will
     * be failed immediately if a content crash occurs. This is also the case when
     * {@link IgnoreCrash} is not present.
     */
    @Target(ElementType.METHOD)
    @Retention(RetentionPolicy.RUNTIME)
    public @interface IgnoreCrash {
        /**
         * @return True if content crashes should be ignored, false otherwise. Default is true.
         */
        boolean value() default true;
    }

    public static class ChildCrashedException extends RuntimeException {
        public ChildCrashedException(final String detailMessage) {
            super(detailMessage);
        }
    }

    public static class RejectedPromiseException extends RuntimeException {
        private final Object mReason;

        /* package */ RejectedPromiseException(final Object reason) {
            super(String.valueOf(reason));
            mReason = reason;
        }

        public Object getReason() {
            return mReason;
        }
    }

    public static class CallRequirement {
        public final boolean allowed;
        public final int count;
        public final int[] order;

        public CallRequirement(final boolean allowed, final int count, final int[] order) {
            this.allowed = allowed;
            this.count = count;
            this.order = order;
        }
    }

    public static class CallInfo {
        public final int counter;
        public final int order;

        /* package */ CallInfo(final int counter, final int order) {
            this.counter = counter;
            this.order = order;
        }
    }

    public static class MethodCall {
        public final GeckoSession session;
        public final Method method;
        public final CallRequirement requirement;
        public final Object target;
        private int currentCount;

        public MethodCall(final GeckoSession session, final Method method,
                          final CallRequirement requirement) {
            this(session, method, requirement, /* target */ null);
        }

        /* package */ MethodCall(final GeckoSession session, final Method method,
                                 final AssertCalled annotation, final Object target) {
            this(session, method,
                 (annotation != null) ? new CallRequirement(annotation.value(),
                                                            annotation.count(),
                                                            annotation.order())
                                      : null,
                 /* target */ target);
        }

        /* package */ MethodCall(final GeckoSession session, final Method method,
                                 final CallRequirement requirement, final Object target) {
            this.session = session;
            this.method = method;
            this.requirement = requirement;
            this.target = target;
            currentCount = 0;
        }

        @Override
        public boolean equals(final Object other) {
            if (this == other) {
                return true;
            } else if (other instanceof MethodCall) {
                final MethodCall otherCall = (MethodCall) other;
                return (session == null || otherCall.session == null ||
                        session.equals(otherCall.session)) &&
                        methodsEqual(method, ((MethodCall) other).method);
            } else if (other instanceof Method) {
                return methodsEqual(method, (Method) other);
            }
            return false;
        }

        @Override
        public int hashCode() {
            return method.hashCode();
        }

        /* package */ int getOrder() {
            if (requirement == null || currentCount == 0) {
                return 0;
            }

            final int[] order = requirement.order;
            if (order == null || order.length == 0) {
                return 0;
            }
            return order[Math.min(currentCount - 1, order.length - 1)];
        }

        /* package */ int getCount() {
            return (requirement == null) ? -1 :
                   requirement.allowed ? requirement.count : 0;
        }

        /* package */ void incrementCounter() {
            currentCount++;
        }

        /* package */ int getCurrentCount() {
            return currentCount;
        }

        /* package */ boolean allowUnlimitedCalls() {
            return getCount() == -1;
        }

        /* package */ boolean allowMoreCalls() {
            final int count = getCount();
            return count == -1 || count > currentCount;
        }

        /* package */ CallInfo getInfo() {
            return new CallInfo(currentCount, getOrder());
        }

        // Similar to Method.equals, but treat the same method from an interface and an
        // overriding class as the same (e.g. CharSequence.length == String.length).
        private static boolean methodsEqual(final @NonNull Method m1, final @NonNull Method m2) {
            return (m1.getDeclaringClass().isAssignableFrom(m2.getDeclaringClass()) ||
                    m2.getDeclaringClass().isAssignableFrom(m1.getDeclaringClass())) &&
                    m1.getName().equals(m2.getName()) &&
                    m1.getReturnType().equals(m2.getReturnType()) &&
                    Arrays.equals(m1.getParameterTypes(), m2.getParameterTypes());
        }
    }

    protected static class CallRecord {
        public final Method method;
        public final MethodCall methodCall;
        public final Object[] args;

        public CallRecord(final GeckoSession session, final Method method, final Object[] args) {
            this.method = method;
            this.methodCall = new MethodCall(session, method, /* requirement */ null);
            this.args = args;
        }
    }

    protected interface CallRecordHandler {
        boolean handleCall(Method method, Object[] args);
    }

    protected final class ExternalDelegate<T> {
        public final Class<T> delegate;
        private final DelegateRegistrar<T> mRegister;
        private final DelegateRegistrar<T> mUnregister;
        private final T mProxy;
        private boolean mRegistered;

        public ExternalDelegate(final Class<T> delegate, final T impl,
                                final DelegateRegistrar<T> register,
                                final DelegateRegistrar<T> unregister) {
            this.delegate = delegate;
            mRegister = register;
            mUnregister = unregister;

            @SuppressWarnings("unchecked")
            final T delegateProxy = (T) Proxy.newProxyInstance(
                    getClass().getClassLoader(), impl.getClass().getInterfaces(),
                    Proxy.getInvocationHandler(mCallbackProxy));
            mProxy = delegateProxy;
        }

        @Override
        public int hashCode() {
            return delegate.hashCode();
        }

        @Override
        public boolean equals(final Object obj) {
            return obj instanceof ExternalDelegate<?> &&
                    delegate.equals(((ExternalDelegate<?>) obj).delegate);
        }

        public void register() {
            try {
                if (!mRegistered) {
                    mRegister.invoke(mProxy);
                    mRegistered = true;
                }
            } catch (final Throwable e) {
                throw unwrapRuntimeException(e);
            }
        }

        public void unregister() {
            try {
                if (mRegistered) {
                    mUnregister.invoke(mProxy);
                    mRegistered = false;
                }
            } catch (final Throwable e) {
                throw unwrapRuntimeException(e);
            }
        }
    }

    protected class CallbackDelegates {
        private final Map<Pair<GeckoSession, Method>, MethodCall> mDelegates = new HashMap<>();
        private final List<ExternalDelegate<?>> mExternalDelegates = new ArrayList<>();
        private int mOrder;
        private JSONObject mOldPrefs;

        public void delegate(final @Nullable GeckoSession session,
                             final @NonNull Object callback) {
            for (final Class<?> ifce : mAllDelegates) {
                if (!ifce.isInstance(callback)) {
                    continue;
                }
                assertThat("Cannot delegate null-delegate callbacks",
                           ifce, not(isIn(mNullDelegates)));
                addDelegatesForInterface(session, callback, ifce);
            }
        }

        private void addDelegatesForInterface(@Nullable final GeckoSession session,
                                              @NonNull final Object callback,
                                              @NonNull final Class<?> ifce) {
            for (final Method method : ifce.getMethods()) {
                final Method callbackMethod;
                try {
                    callbackMethod = callback.getClass().getMethod(method.getName(),
                                                                   method.getParameterTypes());
                } catch (final NoSuchMethodException e) {
                    throw new RuntimeException(e);
                }
                final Pair<GeckoSession, Method> pair = new Pair<>(session, method);
                final MethodCall call = new MethodCall(
                        session, callbackMethod,
                        getAssertCalled(callbackMethod, callback), callback);
                // It's unclear if we should assert the call count if we replace an existing
                // delegate half way through. Until that is resolved, forbid replacing an
                // existing delegate during a test. If you are thinking about changing this
                // behavior, first see if #delegateDuringNextWait fits your needs.
                assertThat("Cannot replace an existing delegate",
                           mDelegates, not(hasKey(pair)));
                mDelegates.put(pair, call);
            }
        }

        public <T> ExternalDelegate<T> addExternalDelegate(
                @NonNull final Class<T> delegate,
                @NonNull final DelegateRegistrar<T> register,
                @NonNull final DelegateRegistrar<T> unregister,
                @NonNull final T impl) {
            assertThat("Delegate must be an interface",
                       delegate.isInterface(), equalTo(true));

            // Delegate each interface to the real thing, then register the delegate using our
            // proxy. That way all calls to the delegate are recorded just like our internal
            // delegates.
            addDelegatesForInterface(/* session */ null, impl, delegate);

            final ExternalDelegate<T> externalDelegate =
                    new ExternalDelegate<>(delegate, impl, register, unregister);
            mExternalDelegates.add(externalDelegate);
            mAllDelegates.add(delegate);
            return externalDelegate;
        }

        @NonNull
        public List<ExternalDelegate<?>> getExternalDelegates() {
            return mExternalDelegates;
        }

        /** Generate a JS function to set new prefs and return a set of saved prefs. */
        public void setPrefs(final @NonNull Map<String, ?> prefs) {
            mOldPrefs = (JSONObject) webExtensionApiCall("SetPrefs", args -> {
                final JSONObject existingPrefs = mOldPrefs != null ? mOldPrefs : new JSONObject();

                final JSONObject newPrefs = new JSONObject();
                for (final Map.Entry<String, ?> pref : prefs.entrySet()) {
                    final Object value = pref.getValue();
                    if (value instanceof Boolean || value instanceof Number ||
                            value instanceof CharSequence) {
                        newPrefs.put(pref.getKey(), value);
                    } else {
                        throw new IllegalArgumentException("Unsupported pref value: " + value);
                    }
                }

                args.put("oldPrefs", existingPrefs);
                args.put("newPrefs", newPrefs);
            });
        }

        /** Generate a JS function to set new prefs and reset a set of saved prefs. */
        private void restorePrefs() {
            if (mOldPrefs == null) {
                return;
            }

            webExtensionApiCall("RestorePrefs", args -> {
                args.put("oldPrefs", mOldPrefs);
                mOldPrefs = null;
            });
        }

        public void clear() {
            for (int i = mExternalDelegates.size() - 1; i >= 0; i--) {
                mExternalDelegates.get(i).unregister();
            }
            mExternalDelegates.clear();
            mDelegates.clear();
            mOrder = 0;

            restorePrefs();
        }

        public void clearAndAssert() {
            final Collection<MethodCall> values = mDelegates.values();
            final MethodCall[] valuesArray = values.toArray(new MethodCall[values.size()]);

            clear();

            for (final MethodCall call : valuesArray) {
                assertMatchesCount(call);
            }
        }

        public MethodCall prepareMethodCall(final GeckoSession session, final Method method) {
            MethodCall call = mDelegates.get(new Pair<>(session, method));
            if (call == null && session != null) {
                call = mDelegates.get(new Pair<>((GeckoSession) null, method));
            }
            if (call == null) {
                return null;
            }

            assertAllowMoreCalls(call);
            call.incrementCounter();
            assertOrder(call, mOrder);
            mOrder = Math.max(call.getOrder(), mOrder);
            return call;
        }
    }

    /* package */ static AssertCalled getAssertCalled(final Method method, final Object callback) {
        final AssertCalled annotation = method.getAnnotation(AssertCalled.class);
        if (annotation != null) {
            return annotation;
        }

        // Some Kotlin lambdas have an invoke method that carries the annotation,
        // instead of the interface method carrying the annotation.
        try {
            return callback.getClass().getDeclaredMethod(
                    "invoke", method.getParameterTypes()).getAnnotation(AssertCalled.class);
        } catch (final NoSuchMethodException e) {
            return null;
        }
    }

    private static void addCallbackClasses(final List<Class<?>> list, final Class<?> ifce) {
        if (!Callbacks.class.equals(ifce.getDeclaringClass())) {
            list.add(ifce);
            return;
        }
        final Class<?>[] superIfces = ifce.getInterfaces();
        for (final Class<?> superIfce : superIfces) {
            addCallbackClasses(list, superIfce);
        }
    }

    private static Set<Class<?>> getDefaultDelegates() {
        final Class<?>[] ifces = Callbacks.class.getDeclaredClasses();
        final List<Class<?>> list = new ArrayList<>(ifces.length);

        for (final Class<?> ifce : ifces) {
            addCallbackClasses(list, ifce);
        }

        return new HashSet<>(list);
    }

    private static final Set<Class<?>> DEFAULT_DELEGATES = getDefaultDelegates();

    public final Environment env = new Environment();

    protected final Instrumentation mInstrumentation =
            InstrumentationRegistry.getInstrumentation();
    protected final GeckoSessionSettings mDefaultSettings;
    protected final Set<GeckoSession> mSubSessions = new HashSet<>();

    protected ErrorCollector mErrorCollector;
    protected GeckoSession mMainSession;
    protected Object mCallbackProxy;
    protected Set<Class<?>> mNullDelegates;
    protected Set<Class<?>> mAllDelegates;
    protected List<CallRecord> mCallRecords;
    protected CallRecordHandler mCallRecordHandler;
    protected CallbackDelegates mWaitScopeDelegates;
    protected CallbackDelegates mTestScopeDelegates;
    protected int mLastWaitStart;
    protected int mLastWaitEnd;
    protected MethodCall mCurrentMethodCall;
    protected long mTimeoutMillis;
    protected Point mDisplaySize;
    protected Map<GeckoSession, SurfaceTexture> mDisplayTextures = new HashMap<>();
    protected Map<GeckoSession, Surface> mDisplaySurfaces = new HashMap<>();
    protected Map<GeckoSession, GeckoDisplay> mDisplays = new HashMap<>();
    protected boolean mClosedSession;
    protected boolean mIgnoreCrash;

    public GeckoSessionTestRule() {
        mDefaultSettings = new GeckoSessionSettings.Builder()
                .build();
    }

    /**
     * Set an ErrorCollector for assertion errors, or null to not use one.
     *
     * @param ec ErrorCollector or null.
     */
    public void setErrorCollector(final @Nullable ErrorCollector ec) {
        mErrorCollector = ec;
    }

    /**
     * Get the current ErrorCollector, or null if not using one.
     *
     * @return ErrorCollector or null.
     */
    public @Nullable ErrorCollector getErrorCollector() {
        return mErrorCollector;
    }

    /**
     * Get the current timeout value in milliseconds.
     *
     * @return The current timeout value in milliseconds.
     */
    public long getTimeoutMillis() {
        return mTimeoutMillis;
    }

    /**
     * Assert a condition with junit.Assert or an error collector.
     *
     * @param reason Reason string
     * @param value Value to check
     * @param matcher Matcher for checking the value
     */
    public <T> void checkThat(final String reason, final T value, final Matcher<? super T> matcher) {
        if (mErrorCollector != null) {
            mErrorCollector.checkThat(reason, value, matcher);
        } else {
            assertThat(reason, value, matcher);
        }
    }

    private void assertAllowMoreCalls(final MethodCall call) {
        final int count = call.getCount();
        if (count != -1) {
            checkThat(call.method.getName() + " call count should be within limit",
                      call.getCurrentCount() + 1, lessThanOrEqualTo(count));
        }
    }

    private void assertOrder(final MethodCall call, final int order) {
        final int newOrder = call.getOrder();
        if (newOrder != 0) {
            checkThat(call.method.getName() + " should be in order",
                      newOrder, greaterThanOrEqualTo(order));
        }
    }

    private void assertMatchesCount(final MethodCall call) {
        if (call.requirement == null) {
            return;
        }
        final int count = call.getCount();
        if (count == 0) {
            checkThat(call.method.getName() + " should not be called",
                      call.getCurrentCount(), equalTo(0));
        } else if (count == -1) {
            checkThat(call.method.getName() + " should be called",
                      call.getCurrentCount(), greaterThan(0));
        } else {
            checkThat(call.method.getName() + " should be called specified number of times",
                      call.getCurrentCount(), equalTo(count));
        }
    }

    /**
     * Get the session set up for the current test.
     *
     * @return GeckoSession object.
     */
    public @NonNull GeckoSession getSession() {
        return mMainSession;
    }

    /**
     * Get the runtime set up for the current test.
     *
     * @return GeckoRuntime object.
     */
    public @NonNull GeckoRuntime getRuntime() {
        return RuntimeCreator.getRuntime();
    }

    public void setTelemetryDelegate(final RuntimeTelemetry.Delegate delegate) {
        RuntimeCreator.setTelemetryDelegate(delegate);
    }

    public @Nullable GeckoDisplay getDisplay() {
        return mDisplays.get(mMainSession);
    }

    protected static Object setDelegate(final @NonNull Class<?> cls,
                                        final @NonNull GeckoSession session,
                                        final @Nullable Object delegate)
            throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
        if (cls == GeckoSession.TextInputDelegate.class) {
            return SessionTextInput.class.getMethod("setDelegate", cls)
                   .invoke(session.getTextInput(), delegate);
        }
        if (cls == ContentBlocking.Delegate.class) {
            return GeckoSession.class.getMethod("setContentBlockingDelegate", cls)
                   .invoke(session, delegate);
        }
        if (cls == Autofill.Delegate.class) {
            return GeckoSession.class.getMethod("setAutofillDelegate", cls)
                   .invoke(session, delegate);
        }
        if (cls == MediaSession.Delegate.class) {
            return GeckoSession.class.getMethod("setMediaSessionDelegate", cls)
                   .invoke(session, delegate);
        }
        return GeckoSession.class.getMethod("set" + cls.getSimpleName(), cls)
               .invoke(session, delegate);
    }

    protected static Object getDelegate(final @NonNull Class<?> cls,
                                        final @NonNull GeckoSession session)
            throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
        if (cls == GeckoSession.TextInputDelegate.class) {
            return SessionTextInput.class.getMethod("getDelegate")
                                         .invoke(session.getTextInput());
        }
        if (cls == ContentBlocking.Delegate.class) {
            return GeckoSession.class.getMethod("getContentBlockingDelegate")
                   .invoke(session);
        }
        if (cls == Autofill.Delegate.class) {
            return GeckoSession.class.getMethod("getAutofillDelegate")
                   .invoke(session);
        }
        if (cls == MediaSession.Delegate.class) {
            return GeckoSession.class.getMethod("getMediaSessionDelegate")
                   .invoke(session);
        }
        return GeckoSession.class.getMethod("get" + cls.getSimpleName())
               .invoke(session);
    }

    @NonNull
    private Set<Class<?>> getCurrentDelegates() {
        final List<ExternalDelegate<?>> waitDelegates = mWaitScopeDelegates.getExternalDelegates();
        final List<ExternalDelegate<?>> testDelegates = mTestScopeDelegates.getExternalDelegates();

        if (waitDelegates.isEmpty() && testDelegates.isEmpty()) {
            return DEFAULT_DELEGATES;
        }

        final Set<Class<?>> set = new HashSet<>(DEFAULT_DELEGATES);
        for (final ExternalDelegate<?> delegate : waitDelegates) {
            set.add(delegate.delegate);
        }
        for (final ExternalDelegate<?> delegate : testDelegates) {
            set.add(delegate.delegate);
        }
        return set;
    }

    private void addNullDelegate(final Class<?> delegate) {
        if (!Callbacks.class.equals(delegate.getDeclaringClass())) {
            assertThat("Null-delegate must be valid interface class",
                       delegate, isIn(DEFAULT_DELEGATES));
            mNullDelegates.add(delegate);
            return;
        }
        for (final Class<?> ifce : delegate.getInterfaces()) {
            addNullDelegate(ifce);
        }
    }

    protected void applyAnnotations(final Collection<Annotation> annotations,
                                    final GeckoSessionSettings settings) {
        for (final Annotation annotation : annotations) {
            if (TimeoutMillis.class.equals(annotation.annotationType())) {
                // Scale timeout based on the default timeout to account for the device under test.
                final long value = ((TimeoutMillis) annotation).value();
                final long timeout = value * env.getScaledTimeoutMillis() / Environment.DEFAULT_TIMEOUT_MILLIS;
                mTimeoutMillis = Math.max(timeout, 1000);
            } else if (Setting.class.equals(annotation.annotationType())) {
                ((Setting) annotation).key().set(settings, ((Setting) annotation).value());
            } else if (Setting.List.class.equals(annotation.annotationType())) {
                for (final Setting setting : ((Setting.List) annotation).value()) {
                    setting.key().set(settings, setting.value());
                }
            } else if (NullDelegate.class.equals(annotation.annotationType())) {
                addNullDelegate(((NullDelegate) annotation).value());
            } else if (NullDelegate.List.class.equals(annotation.annotationType())) {
                for (final NullDelegate nullDelegate : ((NullDelegate.List) annotation).value()) {
                    addNullDelegate(nullDelegate.value());
                }
            } else if (WithDisplay.class.equals(annotation.annotationType())) {
                final WithDisplay displaySize = (WithDisplay)annotation;
                mDisplaySize = new Point(displaySize.width(), displaySize.height());
            } else if (ClosedSessionAtStart.class.equals(annotation.annotationType())) {
                mClosedSession = ((ClosedSessionAtStart) annotation).value();
            } else if (IgnoreCrash.class.equals(annotation.annotationType())) {
                mIgnoreCrash = ((IgnoreCrash) annotation).value();
            }
        }
    }

    private static RuntimeException unwrapRuntimeException(final Throwable e) {
        final Throwable cause = e.getCause();
        if (cause != null && cause instanceof RuntimeException) {
            return (RuntimeException) cause;
        } else if (e instanceof RuntimeException) {
            return (RuntimeException) e;
        }

        return new RuntimeException(cause != null ? cause : e);
    }

    protected void prepareStatement(final Description description) {
        final GeckoSessionSettings settings = new GeckoSessionSettings(mDefaultSettings);
        mTimeoutMillis = env.getDefaultTimeoutMillis();
        mNullDelegates = new HashSet<>();
        mClosedSession = false;
        mIgnoreCrash = false;

        applyAnnotations(Arrays.asList(description.getTestClass().getAnnotations()), settings);
        applyAnnotations(description.getAnnotations(), settings);

        final List<CallRecord> records = new ArrayList<>();
        final CallbackDelegates waitDelegates = new CallbackDelegates();
        final CallbackDelegates testDelegates = new CallbackDelegates();
        mCallRecords = records;
        mWaitScopeDelegates = waitDelegates;
        mTestScopeDelegates = testDelegates;
        mLastWaitStart = 0;
        mLastWaitEnd = 0;

        final InvocationHandler recorder = new InvocationHandler() {
            @Override
            public Object invoke(final Object proxy, final Method method,
                                 final Object[] args) {
                boolean ignore = false;
                MethodCall call = null;

                if (Object.class.equals(method.getDeclaringClass())) {
                    switch (method.getName()) {
                        case "equals":
                            return proxy == args[0];
                        case "toString":
                            return "Call Recorder";
                    }
                    ignore = true;
                } else if (mCallRecordHandler != null) {
                    ignore = mCallRecordHandler.handleCall(method, args);
                }

                final boolean isExternalDelegate =
                        !DEFAULT_DELEGATES.contains(method.getDeclaringClass());

                if (!ignore) {
                    if (!isExternalDelegate) {
                        ThreadUtils.assertOnUiThread();
                    }

                    final GeckoSession session;
                    if (isExternalDelegate) {
                        session = null;
                    } else {
                        assertThat("Callback first argument must be session object",
                                   args, arrayWithSize(greaterThan(0)));
                        assertThat("Callback first argument must be session object",
                                   args[0], instanceOf(GeckoSession.class));
                        session = (GeckoSession) args[0];
                    }

                    if ((sOnCrash.equals(method) || sOnKill.equals(method))
                            && !mIgnoreCrash && isUsingSession(session)) {
                        if (env.shouldShutdownOnCrash()) {
                            getRuntime().shutdown();
                        }

                        throw new ChildCrashedException("Child process crashed");
                    }

                    records.add(new CallRecord(session, method, args));

                    call = waitDelegates.prepareMethodCall(session, method);
                    if (call == null) {
                        call = testDelegates.prepareMethodCall(session, method);
                    }

                    if (isExternalDelegate) {
                        assertThat("External delegate should be registered",
                                   call, notNullValue());
                    }
                }

                Object returnValue = null;
                try {
                    mCurrentMethodCall = call;
                    returnValue = method.invoke((call != null) ? call.target
                                                        : Callbacks.Default.INSTANCE, args);
                } catch (final IllegalAccessException | InvocationTargetException e) {
                    throw unwrapRuntimeException(e);
                } finally {
                    mCurrentMethodCall = null;
                }

                return returnValue;
            }
        };

        final Class<?>[] classes = DEFAULT_DELEGATES.toArray(
                new Class<?>[DEFAULT_DELEGATES.size()]);
        mCallbackProxy = Proxy.newProxyInstance(GeckoSession.class.getClassLoader(),
                                                classes, recorder);
        mAllDelegates = new HashSet<>(DEFAULT_DELEGATES);

        mMainSession = new GeckoSession(settings);
        prepareSession(mMainSession);

        if (mDisplaySize != null) {
            addDisplay(mMainSession, mDisplaySize.x, mDisplaySize.y);
        }

        if (!mClosedSession) {
            openSession(mMainSession);
            UiThreadUtils.waitForCondition(() ->
                            RuntimeCreator.sTestSupport.get() != RuntimeCreator.TEST_SUPPORT_INITIAL,
                    env.getDefaultTimeoutMillis());
            if (RuntimeCreator.sTestSupport.get() != RuntimeCreator.TEST_SUPPORT_OK) {
                throw new RuntimeException("Could not register TestSupport, see logs for error.");
            }
        }
    }

    protected void prepareSession(final GeckoSession session) {
        UiThreadUtils.waitForCondition(() ->
                        RuntimeCreator.sTestSupport.get() != RuntimeCreator.TEST_SUPPORT_INITIAL,
                env.getDefaultTimeoutMillis());
        session.getWebExtensionController()
                .setMessageDelegate(RuntimeCreator.sTestSupportExtension,
                                    mMessageDelegate,
                          "browser");
        for (final Class<?> cls : DEFAULT_DELEGATES) {
            try {
                setDelegate(cls, session, mNullDelegates.contains(cls) ? null : mCallbackProxy);
            } catch (final NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * Call open() on a session, and ensure it's ready for use by the test. In particular,
     * remove any extra calls recorded as part of opening the session.
     *
     * @param session Session to open.
     */
    public void openSession(final GeckoSession session) {
        ThreadUtils.assertOnUiThread();
        // We receive an initial about:blank load; don't expose that to the test. The initial
        // load ends with the first onPageStop call, so ignore everything from the session
        // until the first onPageStop call.

        try {
            // We cannot detect initial page load without progress delegate.
            assertThat("ProgressDelegate cannot be null-delegate when opening session",
                    GeckoSession.ProgressDelegate.class, not(isIn(mNullDelegates)));
            mCallRecordHandler = (method, args) -> {
                Log.e(LOGTAG, "method: " + method);
                final boolean matching = DEFAULT_DELEGATES.contains(
                        method.getDeclaringClass()) && session.equals(args[0]);
                if (matching && sOnPageStop.equals(method)) {
                    mCallRecordHandler = null;
                }
                return matching;
            };

            session.open(getRuntime());

            UiThreadUtils.waitForCondition(() -> mCallRecordHandler == null,
                    env.getDefaultTimeoutMillis());
        } finally {
            mCallRecordHandler = null;
        }
    }

    private void waitForOpenSession(final GeckoSession session) {
        ThreadUtils.assertOnUiThread();
        // We receive an initial about:blank load; don't expose that to the test. The initial
        // load ends with the first onPageStop call, so ignore everything from the session
        // until the first onPageStop call.

        try {
            // We cannot detect initial page load without progress delegate.
            assertThat("ProgressDelegate cannot be null-delegate when opening session",
                       GeckoSession.ProgressDelegate.class, not(isIn(mNullDelegates)));
            mCallRecordHandler = (method, args) -> {
                Log.e(LOGTAG, "method: " + method);
                final boolean matching = DEFAULT_DELEGATES.contains(
                        method.getDeclaringClass()) && session.equals(args[0]);
                if (matching && sOnPageStop.equals(method)) {
                    mCallRecordHandler = null;
                }
                return matching;
            };

            UiThreadUtils.waitForCondition(() -> mCallRecordHandler == null,
                    env.getDefaultTimeoutMillis());
        } finally {
            mCallRecordHandler = null;
        }
    }

    /**
     * Internal method to perform callback checks at the end of a test.
     */
    public void performTestEndCheck() {
        mWaitScopeDelegates.clearAndAssert();
        mTestScopeDelegates.clearAndAssert();
    }

    protected void cleanupSession(final GeckoSession session) {
        if (session.isOpen()) {
            session.close();
        }
        releaseDisplay(session);
    }

    protected boolean isUsingSession(final GeckoSession session) {
        return session.equals(mMainSession) || mSubSessions.contains(session);
    }

    protected void deleteCrashDumps() {
        final File dumpDir = new File(getProfilePath(), "minidumps");
        for (final File dump : dumpDir.listFiles()) {
            dump.delete();
        }
    }

    protected void cleanupExtensions() throws Throwable {
        final WebExtensionController controller = getRuntime().getWebExtensionController();
        final List<WebExtension> list = waitForResult(controller.list());

        boolean hasTestSupport = false;
        // Uninstall any left-over extensions
        for (final WebExtension extension : list) {
            if (!extension.id.equals(RuntimeCreator.TEST_SUPPORT_EXTENSION_ID)) {
                waitForResult(controller.uninstall(extension));
            } else {
                hasTestSupport = true;
            }
        }

        // If an extension was still installed, this test should fail.
        // Note the test support extension is always kept for speed.
        assertThat("A WebExtension was left installed during this test.",
                list.size(), equalTo(hasTestSupport ? 1 : 0));
    }

    protected void cleanupStatement() throws Throwable {
        mWaitScopeDelegates.clear();
        mTestScopeDelegates.clear();

        for (final GeckoSession session : mSubSessions) {
            cleanupSession(session);
        }

        cleanupSession(mMainSession);
        cleanupExtensions();

        if (mIgnoreCrash) {
            deleteCrashDumps();
        }

        mMainSession = null;
        mCallbackProxy = null;
        mAllDelegates = null;
        mNullDelegates = null;
        mCallRecords = null;
        mWaitScopeDelegates = null;
        mTestScopeDelegates = null;
        mLastWaitStart = 0;
        mLastWaitEnd = 0;
        mTimeoutMillis = 0;
        RuntimeCreator.setTelemetryDelegate(null);
    }

    // These markers are used by runjunit.py to capture the logcat of a test
    private static final String TEST_START_MARKER = "test_start 1f0befec-3ff2-40ff-89cf-b127eb38b1ec";
    private static final String TEST_END_MARKER = "test_end c5ee677f-bc83-49bd-9e28-2d35f3d0f059";

    @Override
    public Statement apply(final Statement base, final Description description) {
        return new Statement() {
            private TestServer mServer;

            private void initTest() {
                try {
                    mServer.start(TEST_PORT);

                    RuntimeCreator.setPortDelegate(mPortDelegate);
                    getRuntime();

                    Log.e(LOGTAG, TEST_START_MARKER + " " + description);
                    Log.e(LOGTAG, "before prepareStatement " + description);
                    prepareStatement(description);
                    Log.e(LOGTAG, "after prepareStatement");
                } catch (final Throwable t) {
                    // Any error here is not related to a specific test
                    throw new TestHarnessException(t);
                }
            }

            @Override
            public void evaluate() throws Throwable {
                final AtomicReference<Throwable> exceptionRef = new AtomicReference<>();

                mServer = new TestServer(InstrumentationRegistry.getInstrumentation().getTargetContext());

                mInstrumentation.runOnMainSync(() -> {
                    try {
                        initTest();
                        base.evaluate();
                        Log.e(LOGTAG, "after evaluate");
                        performTestEndCheck();
                        Log.e(LOGTAG, "after performTestEndCheck");
                    } catch (final Throwable t) {
                        Log.e(LOGTAG, "Error", t);
                        exceptionRef.set(t);
                    } finally {
                        try {
                            mServer.stop();
                            cleanupStatement();
                        } catch (final Throwable t) {
                            exceptionRef.compareAndSet(null, t);
                        }
                        Log.e(LOGTAG, TEST_END_MARKER + " " + description);
                    }
                });

                final Throwable throwable = exceptionRef.get();
                if (throwable != null) {
                    throw throwable;
                }
            }
        };
    }

    /**
     * This simply sends an empty message to the web content and waits for a reply.
     */
    public void waitForRoundTrip(final GeckoSession session) {
        waitForJS(session, "true");
    }

    /**
     * Wait until a page load has finished on any session. A session must have started a
     * page load since the last wait, or this method will wait indefinitely.
     */
    public void waitForPageStop() {
        waitForPageStop(/* session */ null);
    }

    /**
     * Wait until a page load has finished. The session must have started a page load since
     * the last wait, or this method will wait indefinitely.
     *
     * @param session Session to wait on, or null to wait on any session.
     */
    public void waitForPageStop(final GeckoSession session) {
        waitForPageStops(session, /* count */ 1);
    }

    /**
     * Wait until a page load has finished on any session. A session must have started a
     * page load since the last wait, or this method will wait indefinitely.
     *
     * @param count Number of page loads to wait for.
     */
    public void waitForPageStops(final int count) {
        waitForPageStops(/* session */ null, count);
    }

    /**
     * Wait until a page load has finished. The session must have started a page load since
     * the last wait, or this method will wait indefinitely.
     *
     * @param session Session to wait on, or null to wait on any session.
     * @param count Number of page loads to wait for.
     */
    public void waitForPageStops(final GeckoSession session, final int count) {
        final List<MethodCall> methodCalls = new ArrayList<>(1);
        methodCalls.add(new MethodCall(session, sOnPageStop,
                new CallRequirement(/* allowed */ true, count, null)));

        waitUntilCalled(session, GeckoSession.ProgressDelegate.class, methodCalls);
    }

    /**
     * Wait until the specified methods have been called on the specified callback
     * interface for any session. If no methods are specified, wait until any method has
     * been called.
     *
     * @param callback Target callback interface; must be an interface under GeckoSession.
     * @param methods List of methods to wait on; use empty or null or wait on any method.
     */
    public void waitUntilCalled(final @NonNull KClass<?> callback,
                                final @Nullable String... methods) {
        waitUntilCalled(/* session */ null, callback, methods);
    }

    /**
     * Wait until the specified methods have been called on the specified callback
     * interface. If no methods are specified, wait until any method has been called.
     *
     * @param session Session to wait on, or null to wait on any session.
     * @param callback Target callback interface; must be an interface under GeckoSession.
     * @param methods List of methods to wait on; use empty or null or wait on any method.
     */
    public void waitUntilCalled(final @Nullable GeckoSession session,
                                final @NonNull KClass<?> callback,
                                final @Nullable String... methods) {
        waitUntilCalled(session, JvmClassMappingKt.getJavaClass(callback), methods);
    }

    /**
     * Wait until the specified methods have been called on the specified callback
     * interface for any session. If no methods are specified, wait until any method has
     * been called.
     *
     * @param callback Target callback interface; must be an interface under GeckoSession.
     * @param methods List of methods to wait on; use empty or null or wait on any method.
     */
    public void waitUntilCalled(final @NonNull Class<?> callback,
                                final @Nullable String... methods) {
        waitUntilCalled(/* session */ null, callback, methods);
    }

    /**
     * Wait until the specified methods have been called on the specified callback
     * interface. If no methods are specified, wait until any method has been called.
     *
     * @param session Session to wait on, or null to wait on any session.
     * @param callback Target callback interface; must be an interface under GeckoSession.
     * @param methods List of methods to wait on; use empty or null or wait on any method.
     */
    public void waitUntilCalled(final @Nullable GeckoSession session,
                                final @NonNull Class<?> callback,
                                final @Nullable String... methods) {
        final int length = (methods != null) ? methods.length : 0;
        final Pattern[] patterns = new Pattern[length];
        for (int i = 0; i < length; i++) {
            patterns[i] = Pattern.compile(methods[i]);
        }

        final List<MethodCall> waitMethods = new ArrayList<>();
        boolean isSessionCallback = false;

        for (final Class<?> ifce : getCurrentDelegates()) {
            if (!ifce.isAssignableFrom(callback)) {
                continue;
            }
            for (final Method method : ifce.getMethods()) {
                for (final Pattern pattern : patterns) {
                    if (!pattern.matcher(method.getName()).matches()) {
                        continue;
                    }
                    waitMethods.add(new MethodCall(session, method,
                                                   /* requirement */ null));
                    break;
                }
            }
            isSessionCallback = true;
        }

        assertThat("Delegate should be a GeckoSession delegate " +
                           "or registered external delegate",
                   isSessionCallback, equalTo(true));

        waitUntilCalled(session, callback, waitMethods);
    }

    /**
     * Wait until the specified methods have been called on the specified object for any
     * session, as specified by any {@link AssertCalled @AssertCalled} annotations. If no
     * {@link AssertCalled @AssertCalled} annotations are found, wait until any method has
     * been called. Only methods belonging to a GeckoSession callback are supported.
     *
     * @param callback Target callback object; must implement an interface under GeckoSession.
     */
    public void waitUntilCalled(final @NonNull Object callback) {
        waitUntilCalled(/* session */ null, callback);
    }

    /**
     * Wait until the specified methods have been called on the specified object,
     * as specified by any {@link AssertCalled @AssertCalled} annotations. If no
     * {@link AssertCalled @AssertCalled} annotations are found, wait until any method
     * has been called. Only methods belonging to a GeckoSession callback are supported.
     *
     * @param session Session to wait on, or null to wait on any session.
     * @param callback Target callback object; must implement an interface under GeckoSession.
     */
    public void waitUntilCalled(final @Nullable GeckoSession session,
                                final @NonNull Object callback) {
        if (callback instanceof Class<?>) {
            waitUntilCalled(session, (Class<?>) callback, (String[]) null);
            return;
        }

        final List<MethodCall> methodCalls = new ArrayList<>();
        boolean isSessionCallback = false;

        for (final Class<?> ifce : getCurrentDelegates()) {
            if (!ifce.isInstance(callback)) {
                continue;
            }
            for (final Method method : ifce.getMethods()) {
                final Method callbackMethod;
                try {
                    callbackMethod = callback.getClass().getMethod(method.getName(),
                                                                   method.getParameterTypes());
                } catch (final NoSuchMethodException e) {
                    throw new RuntimeException(e);
                }
                final AssertCalled ac = getAssertCalled(callbackMethod, callback);
                if (ac != null && ac.value() && ac.count() != 0) {
                    methodCalls.add(new MethodCall(session, method,
                                                   ac, /* target */ null));
                }
            }
            isSessionCallback = true;
        }

        assertThat("Delegate should implement a GeckoSession delegate " +
                           "or registered external delegate",
                   isSessionCallback, equalTo(true));

        waitUntilCalled(session, callback.getClass(), methodCalls);
        forCallbacksDuringWait(session, callback);
    }

    private void waitUntilCalled(final @Nullable GeckoSession session,
                                 final @NonNull Class<?> delegate,
                                 final @NonNull List<MethodCall> methodCalls) {
        ThreadUtils.assertOnUiThread();

        if (session != null && !session.equals(mMainSession)) {
            assertThat("Session should be wrapped through wrapSession",
                       session, isIn(mSubSessions));
        }

        // Make sure all handlers are set though #delegateUntilTestEnd or #delegateDuringNextWait,
        // instead of through GeckoSession directly, so that we can still record calls even with
        // custom handlers set.
        for (final Class<?> ifce : DEFAULT_DELEGATES) {
            final Object callback;
            try {
                callback = getDelegate(ifce, session == null ? mMainSession : session);
            } catch (final NoSuchMethodException | IllegalAccessException |
                    InvocationTargetException e) {
                throw unwrapRuntimeException(e);
            }
            if (mNullDelegates.contains(ifce)) {
                // Null-delegates are initially null but are allowed to be any value.
                continue;
            }
            assertThat(ifce.getSimpleName() + " callbacks should be " +
                       "accessed through GeckoSessionTestRule delegate methods",
                       callback, sameInstance(mCallbackProxy));
        }

        if (methodCalls.isEmpty()) {
            // Waiting for any call on `delegate`; make sure it doesn't contain any null-delegates.
            for (final Class<?> ifce : mNullDelegates) {
                assertThat("Cannot wait on null-delegate callbacks",
                           delegate, not(typeCompatibleWith(ifce)));
            }
        } else {
            // Waiting for particular calls; make sure those calls aren't from a null-delegate.
            for (final MethodCall call : methodCalls) {
                assertThat("Cannot wait on null-delegate callbacks",
                           call.method.getDeclaringClass(), not(isIn(mNullDelegates)));
            }
        }

        boolean calledAny = false;
        int index = mLastWaitEnd;
        final long startTime = SystemClock.uptimeMillis();

        beforeWait();

        while (!calledAny || !methodCalls.isEmpty()) {
            final int currentIndex = index;

            // Let's wait for more messages if we reached the end
            UiThreadUtils.waitForCondition(() -> (currentIndex < mCallRecords.size()), mTimeoutMillis);

            if (SystemClock.uptimeMillis() - startTime > mTimeoutMillis) {
                throw new UiThreadUtils.TimeoutException("Timed out after " + mTimeoutMillis + "ms");
            }

            final MethodCall recorded = mCallRecords.get(index).methodCall;
            calledAny |= recorded.method.getDeclaringClass().isAssignableFrom(delegate);
            index++;

            final int i = methodCalls.indexOf(recorded);
            if (i < 0) {
                continue;
            }

            final MethodCall methodCall = methodCalls.get(i);
            methodCall.incrementCounter();
            if (methodCall.allowUnlimitedCalls() || !methodCall.allowMoreCalls()) {
                methodCalls.remove(i);
            }
        }

        afterWait(index);
    }

    protected void beforeWait() {
        mLastWaitStart = mLastWaitEnd;
    }

    protected void afterWait(final int endCallIndex) {
        mLastWaitEnd = endCallIndex;
        mWaitScopeDelegates.clearAndAssert();

        // Register any test-delegates that were not registered due to wait-delegates
        // having precedence.
        for (final ExternalDelegate<?> delegate : mTestScopeDelegates.getExternalDelegates()) {
            delegate.register();
        }
    }

    /**
     * Playback callbacks that were made on all sessions during the previous wait. For any
     * methods annotated with {@link AssertCalled @AssertCalled}, assert that the
     * callbacks satisfy the specified requirements. If no {@link AssertCalled
     * @AssertCalled} annotations are found, assert any method has been called. Only
     * methods belonging to a GeckoSession callback are supported.
     *
     * @param callback Target callback object; must implement one or more interfaces
     *                 under GeckoSession.
     */
    public void forCallbacksDuringWait(final @NonNull Object callback) {
        forCallbacksDuringWait(/* session */ null, callback);
    }

    /**
     * Playback callbacks that were made during the previous wait. For any methods
     * annotated with {@link AssertCalled @AssertCalled}, assert that the callbacks
     * satisfy the specified requirements. If no {@link AssertCalled @AssertCalled}
     * annotations are found, assert any method has been called. Only methods belonging
     * to a GeckoSession callback are supported.
     *
     * @param session  Target session object, or null to playback all sessions.
     * @param callback Target callback object; must implement one or more interfaces
     *                 under GeckoSession.
     */
    public void forCallbacksDuringWait(final @Nullable GeckoSession session,
                                       final @NonNull Object callback) {
        final Method[] declaredMethods = callback.getClass().getDeclaredMethods();
        final List<MethodCall> methodCalls = new ArrayList<>(declaredMethods.length);
        boolean assertingAnyCall = true;
        Class<?> foundNullDelegate = null;

        for (final Class<?> ifce : mAllDelegates) {
            if (!ifce.isInstance(callback)) {
                continue;
            }
            if (mNullDelegates.contains(ifce)) {
                foundNullDelegate = ifce;
            }
            for (final Method method : ifce.getMethods()) {
                final Method callbackMethod;
                try {
                    callbackMethod = callback.getClass().getMethod(method.getName(),
                                                                   method.getParameterTypes());
                } catch (final NoSuchMethodException e) {
                    throw new RuntimeException(e);
                }
                final MethodCall call = new MethodCall(
                        session, callbackMethod, getAssertCalled(callbackMethod, callback),
                        /* target */ null);
                methodCalls.add(call);

                if (call.requirement != null) {
                    if (foundNullDelegate == ifce) {
                        fail("Cannot assert on null-delegate " + ifce.getSimpleName());
                    }
                    assertingAnyCall = false;
                }
            }
        }

        if (assertingAnyCall && foundNullDelegate != null) {
            fail("Cannot assert on null-delegate " + foundNullDelegate.getSimpleName());
        }

        int order = 0;
        boolean calledAny = false;

        for (int index = mLastWaitStart; index < mLastWaitEnd; index++) {
            final CallRecord record = mCallRecords.get(index);
            if (!record.method.getDeclaringClass().isInstance(callback) ||
                    (session != null && DEFAULT_DELEGATES.contains(
                            record.method.getDeclaringClass()) && !session.equals(record.args[0]))) {
                continue;
            }

            final int i = methodCalls.indexOf(record.methodCall);
            checkThat(record.method.getName() + " should be found",
                      i, greaterThanOrEqualTo(0));

            final MethodCall methodCall = methodCalls.get(i);
            assertAllowMoreCalls(methodCall);
            methodCall.incrementCounter();
            assertOrder(methodCall, order);
            order = Math.max(methodCall.getOrder(), order);

            try {
                mCurrentMethodCall = methodCall;
                record.method.invoke(callback, record.args);
            } catch (final IllegalAccessException | InvocationTargetException e) {
                throw unwrapRuntimeException(e);
            } finally {
                mCurrentMethodCall = null;
            }
            calledAny = true;
        }

        for (final MethodCall methodCall : methodCalls) {
            assertMatchesCount(methodCall);
            if (methodCall.requirement != null) {
                calledAny = true;
            }
        }

        checkThat("Should have called one of " +
                  Arrays.toString(callback.getClass().getInterfaces()),
                  calledAny, equalTo(true));
    }

    /**
     * Get information about the current call. Only valid during a {@link
     * #forCallbacksDuringWait}, {@link #delegateDuringNextWait}, or {@link
     * #delegateUntilTestEnd} callback.
     *
     * @return Call information
     */
    public @NonNull CallInfo getCurrentCall() {
        assertThat("Should be in a method call", mCurrentMethodCall, notNullValue());
        return mCurrentMethodCall.getInfo();
    }

    /**
     * Delegate implemented interfaces to the specified callback object for all sessions,
     * for the rest of the test.  Only GeckoSession callback interfaces are supported.
     * Delegates for {@code delegateUntilTestEnd} can be temporarily overridden by
     * delegates for {@link #delegateDuringNextWait}.
     *
     * @param callback Callback object, or null to clear all previously-set delegates.
     */
    public void delegateUntilTestEnd(final @NonNull Object callback) {
        delegateUntilTestEnd(/* session */ null, callback);
    }

    /**
     * Delegate implemented interfaces to the specified callback object, for the rest of the test.
     * Only GeckoSession callback interfaces are supported. Delegates for {@link
     * #delegateUntilTestEnd} can be temporarily overridden by delegates for {@link
     * #delegateDuringNextWait}.
     *
     * @param session Session to target, or null to target all sessions.
     * @param callback Callback object, or null to clear all previously-set delegates.
     */
    public void delegateUntilTestEnd(final @Nullable GeckoSession session,
                                     final @NonNull Object callback) {
        mTestScopeDelegates.delegate(session, callback);
    }

    /**
     * Delegate implemented interfaces to the specified callback object for all sessions,
     * during the next wait.  Only GeckoSession callback interfaces are supported.
     * Delegates for {@code delegateDuringNextWait} can temporarily take precedence over
     * delegates for {@link #delegateUntilTestEnd}.
     *
     * @param callback Callback object, or null to clear all previously-set delegates.
     */
    public void delegateDuringNextWait(final @NonNull Object callback) {
        delegateDuringNextWait(/* session */ null, callback);
    }

    /**
     * Delegate implemented interfaces to the specified callback object, during the next wait.
     * Only GeckoSession callback interfaces are supported. Delegates for {@link
     * #delegateDuringNextWait} can temporarily take precedence over delegates for
     * {@link #delegateUntilTestEnd}.
     *
     * @param session Session to target, or null to target all sessions.
     * @param callback Callback object, or null to clear all previously-set delegates.
     */
    public void delegateDuringNextWait(final @Nullable GeckoSession session,
                                       final @NonNull Object callback) {
        mWaitScopeDelegates.delegate(session, callback);
    }

    /**
     * Synthesize a tap event at the specified location using the main session.
     * The session must have been created with a display.
     *
     * @param session Target session
     * @param x X coordinate
     * @param y Y coordinate
     */
    public void synthesizeTap(final @NonNull GeckoSession session,
                              final int x, final int y) {
        final long downTime = SystemClock.uptimeMillis();
        final MotionEvent down = MotionEvent.obtain(
                downTime, SystemClock.uptimeMillis(), MotionEvent.ACTION_DOWN, x, y, 0);
        session.getPanZoomController().onTouchEvent(down);

        final MotionEvent up = MotionEvent.obtain(
                downTime, SystemClock.uptimeMillis(), MotionEvent.ACTION_UP, x, y, 0);
        session.getPanZoomController().onTouchEvent(up);
    }

    Map<GeckoSession, WebExtension.Port> mPorts = new HashMap<>();

    private WebExtension.MessageDelegate mMessageDelegate = new WebExtension.MessageDelegate() {
        @Override
        public void onConnect(final @NonNull WebExtension.Port port) {
            mPorts.put(port.sender.session, port);
            port.setDelegate(mPortDelegate);
        }
    };

    private WebExtension.PortDelegate mPortDelegate = new WebExtension.PortDelegate() {
        @Override
        public void onPortMessage(@NonNull final Object message, @NonNull final WebExtension.Port port) {
            final JSONObject response = (JSONObject) message;

            final String id;
            try {
                id = response.getString("id");
                final EvalJSResult result = new EvalJSResult();

                final Object exception = response.get("exception");
                if (exception != JSONObject.NULL) {
                    result.exception = exception;
                }

                final Object value = response.get("response");
                if (value != JSONObject.NULL){
                    result.value = value;
                }

                mPendingMessages.put(id, result);
            } catch (final JSONException ex) {
                throw new RuntimeException(ex);
            }
        }

        @Override
        public void onDisconnect(final @NonNull WebExtension.Port port) {
            mPorts.remove(port.sender.session);
        }
    };

    private static class EvalJSResult {
        Object value;
        Object exception;
    }

    Map<String, EvalJSResult> mPendingMessages = new HashMap<>();

    public class ExtensionPromise {
        private UUID mUuid;
        private GeckoSession mSession;

        protected ExtensionPromise(final UUID uuid, final GeckoSession session, final String js) {
            mUuid = uuid;
            mSession = session;
            evaluateJS(
                    session, "this['" + uuid + "'] = " + js + "; true"
            );
        }

        public Object getValue() {
            return evaluateJS(mSession, "this['" + mUuid + "']");
        }
    }

    public ExtensionPromise evaluatePromiseJS(final @NonNull GeckoSession session,
                                              final @NonNull String js) {
        return new ExtensionPromise(UUID.randomUUID(), session, js);
    }

    public Object evaluateExtensionJS(final @NonNull String js) {
        return webExtensionApiCall("Eval", args -> {
            args.put("code", js);
        });
    }

    public Object evaluateJS(final @NonNull GeckoSession session, final @NonNull String js) {
        // Let's make sure we have the port already
        UiThreadUtils.waitForCondition(() -> mPorts.containsKey(session), mTimeoutMillis);

        final JSONObject message = new JSONObject();
        final String id = UUID.randomUUID().toString();
        try {
            message.put("id", id);
            message.put("eval", js);
        } catch (final JSONException ex) {
            throw new RuntimeException(ex);
        }

        mPorts.get(session).postMessage(message);

        return waitForMessage(id);
    }

    public int getSessionPid(final @NonNull GeckoSession session) {
        final Double dblPid = (Double) webExtensionApiCall(session, "GetPidForTab", null);
        return dblPid.intValue();
    }

    public String getProfilePath() {
        return (String) webExtensionApiCall("GetProfilePath", null);
    }

    public int[] getAllSessionPids() {
        final JSONArray jsonPids = (JSONArray) webExtensionApiCall("GetAllBrowserPids", null);
        final int[] pids = new int[jsonPids.length()];
        for (int i = 0; i < jsonPids.length(); i++) {
            try {
                pids[i] = jsonPids.getInt(i);
            } catch (final JSONException e) {
                throw new RuntimeException(e);
            }
        }
        return pids;
    }

    public void killContentProcess(final int pid) {
        webExtensionApiCall("KillContentProcess", args -> {
            args.put("pid", pid);
        });
    }

    public boolean getActive(final @NonNull GeckoSession session) {
        final Boolean isActive = (Boolean)
                webExtensionApiCall(session, "GetActive", null);
        return isActive;
    }

    private Object waitForMessage(final String id) {
        UiThreadUtils.waitForCondition(() -> mPendingMessages.containsKey(id),
                mTimeoutMillis);

        final EvalJSResult result = mPendingMessages.get(id);
        mPendingMessages.remove(id);

        if (result.exception != null) {
            throw new RejectedPromiseException(result.exception);
        }

        if (result.value == null) {
            return null;
        }

        Object value;
        try {
            value = new JSONTokener((String) result.value).nextValue();
        } catch (final JSONException ex) {
            value = result.value;
        }

        if (value instanceof Integer) {
            return ((Integer) value).doubleValue();
        }
        return value;
    }

    /**
     * Initialize and keep track of the specified session within the test rule. The
     * session is automatically cleaned up at the end of the test.
     *
     * @param session Session to keep track of.
     * @return Same session
     */
    public GeckoSession wrapSession(final GeckoSession session) {
        try {
            mSubSessions.add(session);
            prepareSession(session);
        } catch (final Throwable e) {
            throw unwrapRuntimeException(e);
        }
        return session;
    }

    private GeckoSession createSession(final GeckoSessionSettings settings,
                                       final boolean open) {
        final GeckoSession session = wrapSession(new GeckoSession(settings));
        if (open) {
            openSession(session);
        }
        return session;
    }

    /**
     * Create a new, opened session using the main session settings.
     *
     * @return New session.
     */
    public GeckoSession createOpenSession() {
        return createSession(mMainSession.getSettings(), /* open */ true);
    }

    /**
     * Create a new, opened session using the specified settings.
     *
     * @param settings Settings for the new session.
     * @return New session.
     */
    public GeckoSession createOpenSession(final GeckoSessionSettings settings) {
        return createSession(settings, /* open */ true);
    }

    /**
     * Create a new, closed session using the specified settings.
     *
     * @return New session.
     */
    public GeckoSession createClosedSession() {
        return createSession(mMainSession.getSettings(), /* open */ false);
    }

    /**
     * Create a new, closed session using the specified settings.
     *
     * @param settings Settings for the new session.
     * @return New session.
     */
    public GeckoSession createClosedSession(final GeckoSessionSettings settings) {
        return createSession(settings, /* open */ false);
    }

    /**
     * Return a value from the given array indexed by the current call counter. Only valid
     * during a {@link #forCallbacksDuringWait}, {@link #delegateDuringNextWait}, or
     * {@link #delegateUntilTestEnd} callback.
     * <p><p>
     * Asserts that {@code foo} is equal to {@code "bar"} during the first call and {@code
     * "baz"} during the second call:
     * <pre>{@code assertThat("Foo should match", foo, equalTo(forEachCall("bar",
     * "baz")));}</pre>
     *
     * @param values Input array
     * @return Value from input array indexed by the current call counter.
     */
    @SafeVarargs
    public final <T> T forEachCall(final T... values) {
        assertThat("Should be in a method call", mCurrentMethodCall, notNullValue());
        return values[Math.min(mCurrentMethodCall.getCurrentCount(), values.length) - 1];
    }

    /**
     * Evaluate a JavaScript expression and return the result, similar to {@link #evaluateJS}.
     * In addition, treat the evaluation as a wait event, which will affect other calls such as
     * {@link #forCallbacksDuringWait}. If the result is a Promise, wait on the Promise to settle
     * and return or throw based on the outcome.
     *
     * @param session Session containing the target page.
     * @param js JavaScript expression.
     * @return Result of the expression or value of the resolved Promise.
     * @see #evaluateJS
     */
    public @Nullable Object waitForJS(final @NonNull GeckoSession session, final @NonNull String js) {
        try {
            beforeWait();
            return evaluateJS(session, js);
        } finally {
            afterWait(mCallRecords.size());
        }
    }

    /**
     * Get a list of Gecko prefs. Undefined prefs will return as null.
     *
     * @param prefs List of pref names.
     * @return Pref values as a list of values.
     */
    public JSONArray getPrefs(final @NonNull String... prefs) {
        return (JSONArray) webExtensionApiCall("GetPrefs", args -> {
            args.put("prefs", new JSONArray(Arrays.asList(prefs)));
        });
    }

    /**
     * Gets the color of a link for a given URI and selector.
     *
     * @param uri Page where the link is present.
     * @param selector Selector that matches the link
     * @return String representing the color, e.g. rgb(0, 0, 255)
     */
    public String getLinkColor(final String uri, final String selector) {
        return (String) webExtensionApiCall("GetLinkColor", args -> {
            args.put("uri", uri);
            args.put("selector", selector);
        });
    }

    public List<String> getRequestedLocales() {
        try {
            final JSONArray locales = (JSONArray) webExtensionApiCall("GetRequestedLocales", null);
            final List<String> result = new ArrayList<>();

            for (int i = 0; i < locales.length(); i++) {
                result.add(locales.getString(i));
            }

            return result;
        } catch (final JSONException ex) {
            throw new RuntimeException(ex);
        }
    }

    /**
     * Adds value to the given histogram.
     *
     * @param id the histogram id to increment.
     * @param value to add to the histogram.
     */
    public void addHistogram(final String id, final long value) {
        webExtensionApiCall("AddHistogram", args -> {
            args.put("id", id);
            args.put("value", value);
        });
    }

    /**
     * Revokes all SSL overrides
     */
    public void removeAllCertOverrides() {
        webExtensionApiCall("RemoveAllCertOverrides", null);
    }

    private interface SetArgs {
        void setArgs(JSONObject object) throws JSONException;
    }

    /**
     * Sets value to the given scalar.
     *
     * @param id the scalar to be set.
     * @param value the value to set.
     */
    public <T> void setScalar(final String id, final T value) {
        webExtensionApiCall("SetScalar", args -> {
            args.put("id", id);
            args.put("value", value);
        });
    }

    /**
     * Invokes nsIDOMWindowUtils.setResolutionAndScaleTo.
     */
    public void setResolutionAndScaleTo(final float resolution) {
        webExtensionApiCall("SetResolutionAndScaleTo", args -> {
            args.put("resolution", resolution);
        });
    }

    /**
     * Invokes nsIDOMWindowUtils.flushApzRepaints.
     */
    public void flushApzRepaints(final GeckoSession session) {
        webExtensionApiCall(session, "FlushApzRepaints", null);
    }

    private Object webExtensionApiCall(final @NonNull String apiName, final @NonNull SetArgs argsSetter) {
        return webExtensionApiCall(null, apiName, argsSetter);
    }

    private Object webExtensionApiCall(final GeckoSession session, final @NonNull String apiName,
                                       final @NonNull SetArgs argsSetter) {
        // Ensure background script is connected
        UiThreadUtils.waitForCondition(() -> RuntimeCreator.backgroundPort() != null,
                mTimeoutMillis);

        if (session != null) {
            // Ensure content script is connected
            UiThreadUtils.waitForCondition(() -> mPorts.get(session) != null,
                    mTimeoutMillis);
        }

        final String id = UUID.randomUUID().toString();

        final JSONObject message = new JSONObject();

        try {
            final JSONObject args = new JSONObject();
            if (argsSetter != null) {
                argsSetter.setArgs(args);
            }

            message.put("id", id);
            message.put("type", apiName);
            message.put("args", args);
        } catch (final JSONException ex) {
            throw new RuntimeException(ex);
        }

        if (session == null) {
            RuntimeCreator.backgroundPort().postMessage(message);
        } else {
            // We post the message using session's port instead of the background port. By routing
            // the message through the extension's content script, we are able to obtain and attach
            // the session's WebExtension tab as a `tab` argument to the API.
            mPorts.get(session).postMessage(message);
        }

        return waitForMessage(id);
    }

    /**
     * Set a list of Gecko prefs for the rest of the test. Prefs set in {@link #setPrefsDuringNextWait} can
     * temporarily take precedence over prefs set in {@code setPrefsUntilTestEnd}.
     *
     * @param prefs Map of pref names to values.
     * @see #setPrefsDuringNextWait
     */
    public void setPrefsUntilTestEnd(final @NonNull Map<String, ?> prefs) {
        mTestScopeDelegates.setPrefs(prefs);
    }

    /**
     * Set a list of Gecko prefs during the next wait. Prefs set in {@code setPrefsDuringNextWait} can
     * temporarily take precedence over prefs set in {@link #setPrefsUntilTestEnd}.
     *
     * @param prefs Map of pref names to values.
     * @see #setPrefsUntilTestEnd
     */
    public void setPrefsDuringNextWait(final @NonNull Map<String, ?> prefs) {
        mWaitScopeDelegates.setPrefs(prefs);
    }

    /**
     * Register an external, non-GeckoSession delegate, and start recording the delegate calls
     * until the end of the test. The delegate can then be used with methods such as {@link
     * #waitUntilCalled(Class, String...)} and {@link #forCallbacksDuringWait(Object)}. At the
     * end of the test, the delegate is automatically unregistered. Delegates added by {@link
     * #addExternalDelegateDuringNextWait} can temporarily take precedence over delegates added
     * by {@code delegateUntilTestEnd}.
     *
     * @param delegate Delegate instance to register.
     * @param register DelegateRegistrar instance that represents a function to register the
     *                 delegate.
     * @param unregister DelegateRegistrar instance that represents a function to unregister the
     *                   delegate.
     * @param impl Default delegate implementation. Its methods may be annotated with
     *             {@link AssertCalled} annotations to assert expected behavior.
     * @see #addExternalDelegateDuringNextWait
     */
    public <T> void addExternalDelegateUntilTestEnd(@NonNull final Class<T> delegate,
                                                    @NonNull final DelegateRegistrar<T> register,
                                                    @NonNull final DelegateRegistrar<T> unregister,
                                                    @NonNull final T impl) {
        final ExternalDelegate<T> externalDelegate =
                mTestScopeDelegates.addExternalDelegate(delegate, register, unregister, impl);

        // Register if there is not a wait delegate to take precedence over this call.
        if (!mWaitScopeDelegates.getExternalDelegates().contains(externalDelegate)) {
            externalDelegate.register();
        }
    }

    /** @see #addExternalDelegateUntilTestEnd(Class, DelegateRegistrar,
     *                                        DelegateRegistrar, Object) */
    public <T> void addExternalDelegateUntilTestEnd(@NonNull final KClass<T> delegate,
                                                    @NonNull final DelegateRegistrar<T> register,
                                                    @NonNull final DelegateRegistrar<T> unregister,
                                                    @NonNull final T impl) {
        addExternalDelegateUntilTestEnd(JvmClassMappingKt.getJavaClass(delegate),
                                        register, unregister, impl);
    }

    /**
     * Register an external, non-GeckoSession delegate, and start recording the delegate calls
     * during the next wait. The delegate can then be used with methods such as {@link
     * #waitUntilCalled(Class, String...)} and {@link #forCallbacksDuringWait(Object)}. After the
     * next wait, the delegate is automatically unregistered. Delegates added by {@code
     * addExternalDelegateDuringNextWait} can temporarily take precedence over delegates added
     * by {@link #delegateUntilTestEnd}.
     *
     * @param delegate Delegate instance to register.
     * @param register DelegateRegistrar instance that represents a function to register the
     *                 delegate.
     * @param unregister DelegateRegistrar instance that represents a function to unregister the
     *                   delegate.
     * @param impl Default delegate implementation. Its methods may be annotated with
     *             {@link AssertCalled} annotations to assert expected behavior.
     * @see #addExternalDelegateDuringNextWait
     */
    public <T> void addExternalDelegateDuringNextWait(@NonNull final Class<T> delegate,
                                                      @NonNull final DelegateRegistrar<T> register,
                                                      @NonNull final DelegateRegistrar<T> unregister,
                                                      @NonNull final T impl) {
        final ExternalDelegate<T> externalDelegate =
                mWaitScopeDelegates.addExternalDelegate(delegate, register, unregister, impl);

        // Always register because this call always takes precedence, but make sure to unregister
        // any test-delegates first.
        final int index = mTestScopeDelegates.getExternalDelegates().indexOf(externalDelegate);
        if (index >= 0) {
            mTestScopeDelegates.getExternalDelegates().get(index).unregister();
        }
        externalDelegate.register();
    }

    /** @see #addExternalDelegateDuringNextWait(Class, DelegateRegistrar,
     *                                          DelegateRegistrar, Object) */
    public <T> void addExternalDelegateDuringNextWait(@NonNull final KClass<T> delegate,
                                                      @NonNull final DelegateRegistrar<T> register,
                                                      @NonNull final DelegateRegistrar<T> unregister,
                                                      @NonNull final T impl) {
        addExternalDelegateDuringNextWait(JvmClassMappingKt.getJavaClass(delegate),
                                          register, unregister, impl);
    }

    /**
     * This waits for the given result and returns it's value. If
     * the result failed with an exception, it is rethrown.
     *
     * @param result A {@link GeckoResult} instance.
     * @param <T> The type of the value held by the {@link GeckoResult}
     * @return The value of the completed {@link GeckoResult}.
     */
    public <T> T waitForResult(@NonNull final GeckoResult<T> result) throws Throwable {
        beforeWait();
        try {
            return UiThreadUtils.waitForResult(result, mTimeoutMillis);
        } catch (final Throwable e) {
            throw unwrapRuntimeException(e);
        } finally {
            afterWait(mCallRecords.size());
        }
    }
}
