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.content.BroadcastReceiver;
18import android.content.ComponentCallbacks;
19import android.content.ComponentName;
20import android.content.ContentProviderClient;
21import android.content.Context;
22import android.content.ContextWrapper;
23import android.content.Intent;
24import android.content.IntentFilter;
25import android.content.ServiceConnection;
26import android.content.pm.PackageManager;
27import android.content.res.Resources;
28import android.net.Uri;
29import android.os.Handler;
30import android.os.IBinder;
31import android.os.UserHandle;
32import android.provider.Settings;
33import android.util.ArrayMap;
34import android.view.LayoutInflater;
35
36import org.junit.rules.TestRule;
37import org.junit.rules.TestWatcher;
38import org.junit.runner.Description;
39import org.junit.runners.model.Statement;
40
41/**
42 * A ContextWrapper with utilities specifically designed to make Testing easier.
43 *
44 * <ul>
45 * <li>System services can be mocked out with {@link #addMockSystemService}</li>
46 * <li>Service binding can be mocked out with {@link #addMockService}</li>
47 * <li>Resources can be mocked out using {@link #getOrCreateTestableResources()}</li>
48 * <li>Settings support {@link TestableSettingsProvider}</li>
49 * <li>Has support for {@link LeakCheck} for services and receivers</li>
50 * </ul>
51 *
52 * <p>TestableContext should be defined as a rule on your test so it can clean up after itself.
53 * Like the following:</p>
54 * <pre class="prettyprint">
55 * &#064;Rule
56 * private final TestableContext mContext = new TestableContext(InstrumentationRegister.getContext());
57 * </pre>
58 */
59public class TestableContext extends ContextWrapper implements TestRule {
60
61    private final TestableContentResolver mTestableContentResolver;
62    private final TestableSettingsProvider mSettingsProvider;
63
64    private ArrayMap<String, Object> mMockSystemServices;
65    private ArrayMap<ComponentName, IBinder> mMockServices;
66    private ArrayMap<ServiceConnection, ComponentName> mActiveServices;
67
68    private PackageManager mMockPackageManager;
69    private LeakCheck.Tracker mReceiver;
70    private LeakCheck.Tracker mService;
71    private LeakCheck.Tracker mComponent;
72    private TestableResources mTestableResources;
73    private TestablePermissions mTestablePermissions;
74
75    public TestableContext(Context base) {
76        this(base, null);
77    }
78
79    public TestableContext(Context base, LeakCheck check) {
80        super(base);
81        mTestableContentResolver = new TestableContentResolver(base);
82        ContentProviderClient settings = base.getContentResolver()
83                .acquireContentProviderClient(Settings.AUTHORITY);
84        mSettingsProvider = TestableSettingsProvider.getFakeSettingsProvider(settings);
85        mTestableContentResolver.addProvider(Settings.AUTHORITY, mSettingsProvider);
86        mSettingsProvider.clearValuesAndCheck(TestableContext.this);
87        mReceiver = check != null ? check.getTracker("receiver") : null;
88        mService = check != null ? check.getTracker("service") : null;
89        mComponent = check != null ? check.getTracker("component") : null;
90    }
91
92    public void setMockPackageManager(PackageManager mock) {
93        mMockPackageManager = mock;
94    }
95
96    @Override
97    public PackageManager getPackageManager() {
98        if (mMockPackageManager != null) {
99            return mMockPackageManager;
100        }
101        return super.getPackageManager();
102    }
103
104    /**
105     * Makes sure the resources being returned by this TestableContext are a version of
106     * TestableResources.
107     * @see #getResources()
108     */
109    public void ensureTestableResources() {
110        if (mTestableResources == null) {
111            mTestableResources = new TestableResources(super.getResources());
112        }
113    }
114
115    /**
116     * Get (and create if necessary) {@link TestableResources} for this TestableContext.
117     */
118    public TestableResources getOrCreateTestableResources() {
119        ensureTestableResources();
120        return mTestableResources;
121    }
122
123    /**
124     * Returns a Resources instance for the test.
125     *
126     * By default this returns the same resources object that would come from the
127     * {@link ContextWrapper}, but if {@link #ensureTestableResources()} or
128     * {@link #getOrCreateTestableResources()} has been called, it will return resources gotten from
129     * {@link TestableResources}.
130     */
131    @Override
132    public Resources getResources() {
133        return mTestableResources != null ? mTestableResources.getResources()
134                : super.getResources();
135    }
136
137    /**
138     * @see #getSystemService(String)
139     */
140    public <T> void addMockSystemService(Class<T> service, T mock) {
141        addMockSystemService(getSystemServiceName(service), mock);
142    }
143
144    /**
145     * @see #getSystemService(String)
146     */
147    public void addMockSystemService(String name, Object service) {
148        if (mMockSystemServices == null) mMockSystemServices = new ArrayMap<>();
149        mMockSystemServices.put(name, service);
150    }
151
152    /**
153     * If a matching mock service has been added through {@link #addMockSystemService} then
154     * that will be returned, otherwise the real service will be acquired from the base
155     * context.
156     */
157    @Override
158    public Object getSystemService(String name) {
159        if (mMockSystemServices != null && mMockSystemServices.containsKey(name)) {
160            return mMockSystemServices.get(name);
161        }
162        if (name.equals(LAYOUT_INFLATER_SERVICE)) {
163            return getBaseContext().getSystemService(LayoutInflater.class).cloneInContext(this);
164        }
165        return super.getSystemService(name);
166    }
167
168    TestableSettingsProvider getSettingsProvider() {
169        return mSettingsProvider;
170    }
171
172    @Override
173    public TestableContentResolver getContentResolver() {
174        return mTestableContentResolver;
175    }
176
177    /**
178     * Will always return itself for a TestableContext to ensure the testable effects extend
179     * to the application context.
180     */
181    @Override
182    public Context getApplicationContext() {
183        // Return this so its always a TestableContext.
184        return this;
185    }
186
187    @Override
188    public Intent registerReceiver(BroadcastReceiver receiver, IntentFilter filter) {
189        if (mReceiver != null) mReceiver.getLeakInfo(receiver).addAllocation(new Throwable());
190        return super.registerReceiver(receiver, filter);
191    }
192
193    @Override
194    public Intent registerReceiver(BroadcastReceiver receiver, IntentFilter filter,
195            String broadcastPermission, Handler scheduler) {
196        if (mReceiver != null) mReceiver.getLeakInfo(receiver).addAllocation(new Throwable());
197        return super.registerReceiver(receiver, filter, broadcastPermission, scheduler);
198    }
199
200    @Override
201    public Intent registerReceiverAsUser(BroadcastReceiver receiver, UserHandle user,
202            IntentFilter filter, String broadcastPermission, Handler scheduler) {
203        if (mReceiver != null) mReceiver.getLeakInfo(receiver).addAllocation(new Throwable());
204        return super.registerReceiverAsUser(receiver, user, filter, broadcastPermission,
205                scheduler);
206    }
207
208    @Override
209    public void unregisterReceiver(BroadcastReceiver receiver) {
210        if (mReceiver != null) mReceiver.getLeakInfo(receiver).clearAllocations();
211        super.unregisterReceiver(receiver);
212    }
213
214    /**
215     * Adds a mock service to be connected to by a bindService call.
216     * <p>
217     *     Normally a TestableContext will pass through all bind requests to the base context
218     *     but when addMockService has been called for a ComponentName being bound, then
219     *     TestableContext will immediately trigger a {@link ServiceConnection#onServiceConnected}
220     *     with the specified service, and will call {@link ServiceConnection#onServiceDisconnected}
221     *     when the service is unbound.
222     * </p>
223     */
224    public void addMockService(ComponentName component, IBinder service) {
225        if (mMockServices == null) mMockServices = new ArrayMap<>();
226        mMockServices.put(component, service);
227    }
228
229    /**
230     * @see #addMockService(ComponentName, IBinder)
231     */
232    @Override
233    public boolean bindService(Intent service, ServiceConnection conn, int flags) {
234        if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
235        if (checkMocks(service.getComponent(), conn)) return true;
236        return super.bindService(service, conn, flags);
237    }
238
239    /**
240     * @see #addMockService(ComponentName, IBinder)
241     */
242    @Override
243    public boolean bindServiceAsUser(Intent service, ServiceConnection conn, int flags,
244            Handler handler, UserHandle user) {
245        if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
246        if (checkMocks(service.getComponent(), conn)) return true;
247        return super.bindServiceAsUser(service, conn, flags, handler, user);
248    }
249
250    /**
251     * @see #addMockService(ComponentName, IBinder)
252     */
253    @Override
254    public boolean bindServiceAsUser(Intent service, ServiceConnection conn, int flags,
255            UserHandle user) {
256        if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
257        if (checkMocks(service.getComponent(), conn)) return true;
258        return super.bindServiceAsUser(service, conn, flags, user);
259    }
260
261    private boolean checkMocks(ComponentName component, ServiceConnection conn) {
262        if (mMockServices != null && component != null && mMockServices.containsKey(component)) {
263            if (mActiveServices == null) mActiveServices = new ArrayMap<>();
264            mActiveServices.put(conn, component);
265            conn.onServiceConnected(component, mMockServices.get(component));
266            return true;
267        }
268        return false;
269    }
270
271    /**
272     * @see #addMockService(ComponentName, IBinder)
273     */
274    @Override
275    public void unbindService(ServiceConnection conn) {
276        if (mService != null) mService.getLeakInfo(conn).clearAllocations();
277        if (mActiveServices != null && mActiveServices.containsKey(conn)) {
278            conn.onServiceDisconnected(mActiveServices.get(conn));
279            mActiveServices.remove(conn);
280            return;
281        }
282        super.unbindService(conn);
283    }
284
285    /**
286     * Check if the TestableContext has a mock binding for a specified component. Will return
287     * true between {@link ServiceConnection#onServiceConnected} and
288     * {@link ServiceConnection#onServiceDisconnected} callbacks for a mock service.
289     *
290     * @see #addMockService(ComponentName, IBinder)
291     */
292    public boolean isBound(ComponentName component) {
293        return mActiveServices != null && mActiveServices.containsValue(component);
294    }
295
296    @Override
297    public void registerComponentCallbacks(ComponentCallbacks callback) {
298        if (mComponent != null) mComponent.getLeakInfo(callback).addAllocation(new Throwable());
299        super.registerComponentCallbacks(callback);
300    }
301
302    @Override
303    public void unregisterComponentCallbacks(ComponentCallbacks callback) {
304        if (mComponent != null) mComponent.getLeakInfo(callback).clearAllocations();
305        super.unregisterComponentCallbacks(callback);
306    }
307
308    public TestablePermissions getTestablePermissions() {
309        if (mTestablePermissions == null) {
310            mTestablePermissions = new TestablePermissions();
311        }
312        return mTestablePermissions;
313    }
314
315    @Override
316    public int checkCallingOrSelfPermission(String permission) {
317        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
318            return mTestablePermissions.check(permission);
319        }
320        return super.checkCallingOrSelfPermission(permission);
321    }
322
323    @Override
324    public int checkCallingPermission(String permission) {
325        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
326            return mTestablePermissions.check(permission);
327        }
328        return super.checkCallingPermission(permission);
329    }
330
331    @Override
332    public int checkPermission(String permission, int pid, int uid) {
333        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
334            return mTestablePermissions.check(permission);
335        }
336        return super.checkPermission(permission, pid, uid);
337    }
338
339    @Override
340    public int checkPermission(String permission, int pid, int uid, IBinder callerToken) {
341        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
342            return mTestablePermissions.check(permission);
343        }
344        return super.checkPermission(permission, pid, uid, callerToken);
345    }
346
347    @Override
348    public int checkSelfPermission(String permission) {
349        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
350            return mTestablePermissions.check(permission);
351        }
352        return super.checkSelfPermission(permission);
353    }
354
355    @Override
356    public void enforceCallingOrSelfPermission(String permission, String message) {
357        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
358            mTestablePermissions.enforce(permission);
359        } else {
360            super.enforceCallingOrSelfPermission(permission, message);
361        }
362    }
363
364    @Override
365    public void enforceCallingPermission(String permission, String message) {
366        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
367            mTestablePermissions.enforce(permission);
368        } else {
369            super.enforceCallingPermission(permission, message);
370        }
371    }
372
373    @Override
374    public void enforcePermission(String permission, int pid, int uid, String message) {
375        if (mTestablePermissions != null && mTestablePermissions.wantsCall(permission)) {
376            mTestablePermissions.enforce(permission);
377        } else {
378            super.enforcePermission(permission, pid, uid, message);
379        }
380    }
381
382    @Override
383    public int checkCallingOrSelfUriPermission(Uri uri, int modeFlags) {
384        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
385            return mTestablePermissions.check(uri, modeFlags);
386        }
387        return super.checkCallingOrSelfUriPermission(uri, modeFlags);
388    }
389
390    @Override
391    public int checkCallingUriPermission(Uri uri, int modeFlags) {
392        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
393            return mTestablePermissions.check(uri, modeFlags);
394        }
395        return super.checkCallingUriPermission(uri, modeFlags);
396    }
397
398    @Override
399    public void enforceCallingOrSelfUriPermission(Uri uri, int modeFlags, String message) {
400        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
401            mTestablePermissions.enforce(uri, modeFlags);
402        } else {
403            super.enforceCallingOrSelfUriPermission(uri, modeFlags, message);
404        }
405    }
406
407    @Override
408    public int checkUriPermission(Uri uri, int pid, int uid, int modeFlags) {
409        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
410            return mTestablePermissions.check(uri, modeFlags);
411        }
412        return super.checkUriPermission(uri, pid, uid, modeFlags);
413    }
414
415    @Override
416    public int checkUriPermission(Uri uri, int pid, int uid, int modeFlags, IBinder callerToken) {
417        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
418            return mTestablePermissions.check(uri, modeFlags);
419        }
420        return super.checkUriPermission(uri, pid, uid, modeFlags, callerToken);
421    }
422
423    @Override
424    public int checkUriPermission(Uri uri, String readPermission, String writePermission, int pid,
425            int uid, int modeFlags) {
426        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
427            return mTestablePermissions.check(uri, modeFlags);
428        }
429        return super.checkUriPermission(uri, readPermission, writePermission, pid, uid, modeFlags);
430    }
431
432    @Override
433    public void enforceCallingUriPermission(Uri uri, int modeFlags, String message) {
434        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
435            mTestablePermissions.enforce(uri, modeFlags);
436        } else {
437            super.enforceCallingUriPermission(uri, modeFlags, message);
438        }
439    }
440
441    @Override
442    public void enforceUriPermission(Uri uri, int pid, int uid, int modeFlags, String message) {
443        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
444            mTestablePermissions.enforce(uri, modeFlags);
445        } else {
446            super.enforceUriPermission(uri, pid, uid, modeFlags, message);
447        }
448    }
449
450    @Override
451    public void enforceUriPermission(Uri uri, String readPermission, String writePermission,
452            int pid, int uid, int modeFlags, String message) {
453        if (mTestablePermissions != null && mTestablePermissions.wantsCall(uri)) {
454            mTestablePermissions.enforce(uri, modeFlags);
455        } else {
456            super.enforceUriPermission(uri, readPermission, writePermission, pid, uid, modeFlags,
457                    message);
458        }
459    }
460
461    @Override
462    public Statement apply(Statement base, Description description) {
463        return new TestWatcher() {
464            @Override
465            protected void succeeded(Description description) {
466                mSettingsProvider.clearValuesAndCheck(TestableContext.this);
467            }
468
469            @Override
470            protected void failed(Throwable e, Description description) {
471                mSettingsProvider.clearValuesAndCheck(TestableContext.this);
472            }
473        }.apply(base, description);
474    }
475}
476