Home | History | Annotate | Download | only in testing
      1 /*
      2  * Copyright (C) 2017 The Android Open Source Project
      3  *
      4  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
      5  * except in compliance with the License. You may obtain a copy of the License at
      6  *
      7  *      http://www.apache.org/licenses/LICENSE-2.0
      8  *
      9  * Unless required by applicable law or agreed to in writing, software distributed under the
     10  * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
     11  * KIND, either express or implied. See the License for the specific language governing
     12  * permissions and limitations under the License.
     13  */
     14 
     15 package android.testing;
     16 
     17 import android.os.Handler;
     18 import android.os.HandlerThread;
     19 import android.os.Looper;
     20 import android.os.Message;
     21 import android.os.MessageQueue;
     22 import android.os.TestLooperManager;
     23 import android.util.ArrayMap;
     24 
     25 import androidx.test.InstrumentationRegistry;
     26 
     27 import org.junit.runners.model.FrameworkMethod;
     28 
     29 import java.lang.annotation.ElementType;
     30 import java.lang.annotation.Retention;
     31 import java.lang.annotation.RetentionPolicy;
     32 import java.lang.annotation.Target;
     33 import java.util.Map;
     34 
     35 /**
     36  * This is a wrapper around {@link TestLooperManager} to make it easier to manage
     37  * and provide an easy annotation for use with tests.
     38  *
     39  * @see TestableLooperTest TestableLooperTest for examples.
     40  */
     41 public class TestableLooper {
     42 
     43     /**
     44      * Whether to hold onto the main thread through all tests in an attempt to
     45      * catch crashes.
     46      */
     47     public static final boolean HOLD_MAIN_THREAD = false;
     48 
     49     private Looper mLooper;
     50     private MessageQueue mQueue;
     51     private MessageHandler mMessageHandler;
     52 
     53     private Handler mHandler;
     54     private Runnable mEmptyMessage;
     55     private TestLooperManager mQueueWrapper;
     56 
     57     public TestableLooper(Looper l) throws Exception {
     58         this(acquireLooperManager(l), l);
     59     }
     60 
     61     private TestableLooper(TestLooperManager wrapper, Looper l) {
     62         mQueueWrapper = wrapper;
     63         setupQueue(l);
     64     }
     65 
     66     private TestableLooper(Looper looper, boolean b) {
     67         setupQueue(looper);
     68     }
     69 
     70     public Looper getLooper() {
     71         return mLooper;
     72     }
     73 
     74     private void setupQueue(Looper l) {
     75         mLooper = l;
     76         mQueue = mLooper.getQueue();
     77         mHandler = new Handler(mLooper);
     78     }
     79 
     80     /**
     81      * Must be called to release the looper when the test is complete, otherwise
     82      * the looper will not be available for any subsequent tests. This is
     83      * automatically handled for tests using {@link RunWithLooper}.
     84      */
     85     public void destroy() {
     86         mQueueWrapper.release();
     87         if (HOLD_MAIN_THREAD && mLooper == Looper.getMainLooper()) {
     88             TestableInstrumentation.releaseMain();
     89         }
     90     }
     91 
     92     /**
     93      * Sets a callback for all messages processed on this TestableLooper.
     94      *
     95      * @see {@link MessageHandler}
     96      */
     97     public void setMessageHandler(MessageHandler handler) {
     98         mMessageHandler = handler;
     99     }
    100 
    101     /**
    102      * Parse num messages from the message queue.
    103      *
    104      * @param num Number of messages to parse
    105      */
    106     public int processMessages(int num) {
    107         for (int i = 0; i < num; i++) {
    108             if (!parseMessageInt()) {
    109                 return i + 1;
    110             }
    111         }
    112         return num;
    113     }
    114 
    115     /**
    116      * Process messages in the queue until no more are found.
    117      */
    118     public void processAllMessages() {
    119         while (processQueuedMessages() != 0) ;
    120     }
    121 
    122     private int processQueuedMessages() {
    123         int count = 0;
    124         mEmptyMessage = () -> { };
    125         mHandler.post(mEmptyMessage);
    126         waitForMessage(mQueueWrapper, mHandler, mEmptyMessage);
    127         while (parseMessageInt()) count++;
    128         return count;
    129     }
    130 
    131     private boolean parseMessageInt() {
    132         try {
    133             Message result = mQueueWrapper.next();
    134             if (result != null) {
    135                 // This is a break message.
    136                 if (result.getCallback() == mEmptyMessage) {
    137                     mQueueWrapper.recycle(result);
    138                     return false;
    139                 }
    140 
    141                 if (mMessageHandler != null) {
    142                     if (mMessageHandler.onMessageHandled(result)) {
    143                         mQueueWrapper.execute(result);
    144                         mQueueWrapper.recycle(result);
    145                     } else {
    146                         mQueueWrapper.recycle(result);
    147                         // Message handler indicated it doesn't want us to continue.
    148                         return false;
    149                     }
    150                 } else {
    151                     mQueueWrapper.execute(result);
    152                     mQueueWrapper.recycle(result);
    153                 }
    154             } else {
    155                 // No messages, don't continue parsing
    156                 return false;
    157             }
    158         } catch (Exception e) {
    159             throw new RuntimeException(e);
    160         }
    161         return true;
    162     }
    163 
    164     /**
    165      * Runs an executable with myLooper set and processes all messages added.
    166      */
    167     public void runWithLooper(RunnableWithException runnable) throws Exception {
    168         new Handler(getLooper()).post(() -> {
    169             try {
    170                 runnable.run();
    171             } catch (Exception e) {
    172                 throw new RuntimeException(e);
    173             }
    174         });
    175         processAllMessages();
    176     }
    177 
    178     public interface RunnableWithException {
    179         void run() throws Exception;
    180     }
    181 
    182     /**
    183      * Annotation that tells the {@link AndroidTestingRunner} to create a TestableLooper and
    184      * run this test/class on that thread. The {@link TestableLooper} can be acquired using
    185      * {@link #get(Object)}.
    186      */
    187     @Retention(RetentionPolicy.RUNTIME)
    188     @Target({ElementType.METHOD, ElementType.TYPE})
    189     public @interface RunWithLooper {
    190         boolean setAsMainLooper() default false;
    191     }
    192 
    193     private static void waitForMessage(TestLooperManager queueWrapper, Handler handler,
    194             Runnable execute) {
    195         for (int i = 0; i < 10; i++) {
    196             if (!queueWrapper.hasMessages(handler, null, execute)) {
    197                 try {
    198                     Thread.sleep(1);
    199                 } catch (InterruptedException e) {
    200                 }
    201             }
    202         }
    203         if (!queueWrapper.hasMessages(handler, null, execute)) {
    204             throw new RuntimeException("Message didn't queue...");
    205         }
    206     }
    207 
    208     private static TestLooperManager acquireLooperManager(Looper l) {
    209         if (HOLD_MAIN_THREAD && l == Looper.getMainLooper()) {
    210             TestableInstrumentation.acquireMain();
    211         }
    212         return InstrumentationRegistry.getInstrumentation().acquireLooperManager(l);
    213     }
    214 
    215     private static final Map<Object, TestableLooper> sLoopers = new ArrayMap<>();
    216 
    217     /**
    218      * For use with {@link RunWithLooper}, used to get the TestableLooper that was
    219      * automatically created for this test.
    220      */
    221     public static TestableLooper get(Object test) {
    222         return sLoopers.get(test);
    223     }
    224 
    225     static class LooperFrameworkMethod extends FrameworkMethod {
    226         private HandlerThread mHandlerThread;
    227 
    228         private final TestableLooper mTestableLooper;
    229         private final Looper mLooper;
    230         private final Handler mHandler;
    231 
    232         public LooperFrameworkMethod(FrameworkMethod base, boolean setAsMain, Object test) {
    233             super(base.getMethod());
    234             try {
    235                 mLooper = setAsMain ? Looper.getMainLooper() : createLooper();
    236                 mTestableLooper = new TestableLooper(mLooper, false);
    237             } catch (Exception e) {
    238                 throw new RuntimeException(e);
    239             }
    240             sLoopers.put(test, mTestableLooper);
    241             mHandler = new Handler(mLooper);
    242         }
    243 
    244         public LooperFrameworkMethod(TestableLooper other, FrameworkMethod base) {
    245             super(base.getMethod());
    246             mLooper = other.mLooper;
    247             mTestableLooper = other;
    248             mHandler = Handler.createAsync(mLooper);
    249         }
    250 
    251         public static FrameworkMethod get(FrameworkMethod base, boolean setAsMain, Object test) {
    252             if (sLoopers.containsKey(test)) {
    253                 return new LooperFrameworkMethod(sLoopers.get(test), base);
    254             }
    255             return new LooperFrameworkMethod(base, setAsMain, test);
    256         }
    257 
    258         @Override
    259         public Object invokeExplosively(Object target, Object... params) throws Throwable {
    260             if (Looper.myLooper() == mLooper) {
    261                 // Already on the right thread from another statement, just execute then.
    262                 return super.invokeExplosively(target, params);
    263             }
    264             boolean set = mTestableLooper.mQueueWrapper == null;
    265             if (set) {
    266                 mTestableLooper.mQueueWrapper = acquireLooperManager(mLooper);
    267             }
    268             try {
    269                 Object[] ret = new Object[1];
    270                 // Run the execution on the looper thread.
    271                 Runnable execute = () -> {
    272                     try {
    273                         ret[0] = super.invokeExplosively(target, params);
    274                     } catch (Throwable throwable) {
    275                         throw new LooperException(throwable);
    276                     }
    277                 };
    278                 Message m = Message.obtain(mHandler, execute);
    279 
    280                 // Dispatch our message.
    281                 try {
    282                     mTestableLooper.mQueueWrapper.execute(m);
    283                 } catch (LooperException e) {
    284                     throw e.getSource();
    285                 } catch (RuntimeException re) {
    286                     // If the TestLooperManager has to post, it will wrap what it throws in a
    287                     // RuntimeException, make sure we grab the actual source.
    288                     if (re.getCause() instanceof LooperException) {
    289                         throw ((LooperException) re.getCause()).getSource();
    290                     } else {
    291                         throw re.getCause();
    292                     }
    293                 } finally {
    294                     m.recycle();
    295                 }
    296                 return ret[0];
    297             } finally {
    298                 if (set) {
    299                     mTestableLooper.mQueueWrapper.release();
    300                     mTestableLooper.mQueueWrapper = null;
    301                     if (HOLD_MAIN_THREAD && mLooper == Looper.getMainLooper()) {
    302                         TestableInstrumentation.releaseMain();
    303                     }
    304                 }
    305             }
    306         }
    307 
    308         private Looper createLooper() {
    309             // TODO: Find way to share these.
    310             mHandlerThread = new HandlerThread(TestableLooper.class.getSimpleName());
    311             mHandlerThread.start();
    312             return mHandlerThread.getLooper();
    313         }
    314 
    315         @Override
    316         protected void finalize() throws Throwable {
    317             super.finalize();
    318             if (mHandlerThread != null) {
    319                 mHandlerThread.quit();
    320             }
    321         }
    322 
    323         private static class LooperException extends RuntimeException {
    324             private final Throwable mSource;
    325 
    326             public LooperException(Throwable t) {
    327                 mSource = t;
    328             }
    329 
    330             public Throwable getSource() {
    331                 return mSource;
    332             }
    333         }
    334     }
    335 
    336     /**
    337      * Callback to control the execution of messages on the looper, when set with
    338      * {@link #setMessageHandler(MessageHandler)} then {@link #onMessageHandled(Message)}
    339      * will get called back for every message processed on the {@link TestableLooper}.
    340      */
    341     public interface MessageHandler {
    342         /**
    343          * Return true to have the message executed and delivered to target.
    344          * Return false to not execute the message and stop executing messages.
    345          */
    346         boolean onMessageHandled(Message m);
    347     }
    348 }
    349