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
15package android.testing;
16
17import android.os.Handler;
18import android.os.HandlerThread;
19import android.os.Looper;
20import android.os.Message;
21import android.os.MessageQueue;
22import android.os.TestLooperManager;
23import android.support.test.InstrumentationRegistry;
24import android.util.ArrayMap;
25
26import org.junit.runners.model.FrameworkMethod;
27
28import java.lang.annotation.ElementType;
29import java.lang.annotation.Retention;
30import java.lang.annotation.RetentionPolicy;
31import java.lang.annotation.Target;
32import 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 */
40public 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
50    public TestableLooper(Looper l) throws Exception {
51        this(acquireLooperManager(l), l);
52    }
53
54    private TestableLooper(TestLooperManager wrapper, Looper l) {
55        mQueueWrapper = wrapper;
56        setupQueue(l);
57    }
58
59    private TestableLooper(Looper looper, boolean b) {
60        setupQueue(looper);
61    }
62
63    public Looper getLooper() {
64        return mLooper;
65    }
66
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     */
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     */
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     */
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     */
111    public void processAllMessages() {
112        while (processQueuedMessages() != 0) ;
113    }
114
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
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     */
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 {
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 {
183        boolean setAsMainLooper() default false;
184    }
185
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
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     */
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
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
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
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
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
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
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
319            public LooperException(Throwable t) {
320                mSource = t;
321            }
322
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         */
339        boolean onMessageHandled(Message m);
340    }
341}
342