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