rules/src/main/java/android/support/test/internal/statement/UiThreadStatement.java

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

package android.support.test.internal.statement;

import org.junit.runners.model.Statement;

import java.util.concurrent.atomic.AtomicReference;

import static android.support.test.InstrumentationRegistry.getInstrumentation;

/**
 * {@link Statement} that executes a test on the application's main thread (or UI thread).
 */
public class UiThreadStatement extends Statement {
    private final Statement mBase;
    private final boolean mRunOnUiThread;

    public UiThreadStatement(Statement base, boolean runOnUiThread) {
        mBase = base;
        mRunOnUiThread = runOnUiThread;
    }

    @Override
    public void evaluate() throws Throwable {
        if (mRunOnUiThread) {
            final AtomicReference<Throwable> exceptionRef = new AtomicReference<>();
            getInstrumentation().runOnMainSync(new Runnable() {
                public void run() {
                    try {
                        mBase.evaluate();
                    } catch (Throwable throwable) {
                        exceptionRef.set(throwable);
                    }
                }
            });
            Throwable throwable = exceptionRef.get();
            if (throwable != null) {
                throw throwable;
            }
        } else {
            mBase.evaluate();
        }
    }
}