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.os.Handler;
29import android.os.IBinder;
30import android.os.UserHandle;
31import android.provider.Settings;
32import android.util.ArrayMap;
33import android.view.LayoutInflater;
34
35import org.junit.rules.TestRule;
36import org.junit.rules.TestWatcher;
37import org.junit.runner.Description;
38import org.junit.runners.model.Statement;
39
40/**
41 * A ContextWrapper with utilities specifically designed to make Testing easier.
42 *
43 * <ul>
44 * <li>System services can be mocked out with {@link #addMockSystemService}</li>
45 * <li>Service binding can be mocked out with {@link #addMockService}</li>
46 * <li>Settings support {@link TestableSettingsProvider}</li>
47 * <li>Has support for {@link LeakCheck} for services and receivers</li>
48 * </ul>
49 *
50 * <p>TestableContext should be defined as a rule on your test so it can clean up after itself.
51 * Like the following:</p>
52 * <pre class="prettyprint">
53 * {@literal
54 * @Rule
55 * private final TestableContext mContext = new TestableContext(InstrumentationRegister.getContext());
56 * }
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
73    public TestableContext(Context base) {
74        this(base, null);
75    }
76
77    public TestableContext(Context base, LeakCheck check) {
78        super(base);
79        mTestableContentResolver = new TestableContentResolver(base);
80        ContentProviderClient settings = base.getContentResolver()
81                .acquireContentProviderClient(Settings.AUTHORITY);
82        mSettingsProvider = TestableSettingsProvider.getFakeSettingsProvider(settings);
83        mTestableContentResolver.addProvider(Settings.AUTHORITY, mSettingsProvider);
84        mReceiver = check != null ? check.getTracker("receiver") : null;
85        mService = check != null ? check.getTracker("service") : null;
86        mComponent = check != null ? check.getTracker("component") : null;
87    }
88
89    public void setMockPackageManager(PackageManager mock) {
90        mMockPackageManager = mock;
91    }
92
93    @Override
94    public PackageManager getPackageManager() {
95        if (mMockPackageManager != null) {
96            return mMockPackageManager;
97        }
98        return super.getPackageManager();
99    }
100
101    @Override
102    public Resources getResources() {
103        return super.getResources();
104    }
105
106    public <T> void addMockSystemService(Class<T> service, T mock) {
107        addMockSystemService(getSystemServiceName(service), mock);
108    }
109
110    public void addMockSystemService(String name, Object service) {
111        if (mMockSystemServices == null) mMockSystemServices = new ArrayMap<>();
112        mMockSystemServices.put(name, service);
113    }
114
115    public void addMockService(ComponentName component, IBinder service) {
116        if (mMockServices == null) mMockServices = new ArrayMap<>();
117        mMockServices.put(component, service);
118    }
119
120    @Override
121    public Object getSystemService(String name) {
122        if (mMockSystemServices != null && mMockSystemServices.containsKey(name)) {
123            return mMockSystemServices.get(name);
124        }
125        if (name.equals(LAYOUT_INFLATER_SERVICE)) {
126            return getBaseContext().getSystemService(LayoutInflater.class).cloneInContext(this);
127        }
128        return super.getSystemService(name);
129    }
130
131    TestableSettingsProvider getSettingsProvider() {
132        return mSettingsProvider;
133    }
134
135    @Override
136    public TestableContentResolver getContentResolver() {
137        return mTestableContentResolver;
138    }
139
140    @Override
141    public Context getApplicationContext() {
142        // Return this so its always a TestableContext.
143        return this;
144    }
145
146    @Override
147    public Intent registerReceiver(BroadcastReceiver receiver, IntentFilter filter) {
148        if (mReceiver != null) mReceiver.getLeakInfo(receiver).addAllocation(new Throwable());
149        return super.registerReceiver(receiver, filter);
150    }
151
152    @Override
153    public Intent registerReceiver(BroadcastReceiver receiver, IntentFilter filter,
154            String broadcastPermission, Handler scheduler) {
155        if (mReceiver != null) mReceiver.getLeakInfo(receiver).addAllocation(new Throwable());
156        return super.registerReceiver(receiver, filter, broadcastPermission, scheduler);
157    }
158
159    @Override
160    public Intent registerReceiverAsUser(BroadcastReceiver receiver, UserHandle user,
161            IntentFilter filter, String broadcastPermission, Handler scheduler) {
162        if (mReceiver != null) mReceiver.getLeakInfo(receiver).addAllocation(new Throwable());
163        return super.registerReceiverAsUser(receiver, user, filter, broadcastPermission,
164                scheduler);
165    }
166
167    @Override
168    public void unregisterReceiver(BroadcastReceiver receiver) {
169        if (mReceiver != null) mReceiver.getLeakInfo(receiver).clearAllocations();
170        super.unregisterReceiver(receiver);
171    }
172
173    @Override
174    public boolean bindService(Intent service, ServiceConnection conn, int flags) {
175        if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
176        if (checkMocks(service.getComponent(), conn)) return true;
177        return super.bindService(service, conn, flags);
178    }
179
180    @Override
181    public boolean bindServiceAsUser(Intent service, ServiceConnection conn, int flags,
182            Handler handler, UserHandle user) {
183        if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
184        if (checkMocks(service.getComponent(), conn)) return true;
185        return super.bindServiceAsUser(service, conn, flags, handler, user);
186    }
187
188    @Override
189    public boolean bindServiceAsUser(Intent service, ServiceConnection conn, int flags,
190            UserHandle user) {
191        if (mService != null) mService.getLeakInfo(conn).addAllocation(new Throwable());
192        if (checkMocks(service.getComponent(), conn)) return true;
193        return super.bindServiceAsUser(service, conn, flags, user);
194    }
195
196    private boolean checkMocks(ComponentName component, ServiceConnection conn) {
197        if (mMockServices != null && component != null && mMockServices.containsKey(component)) {
198            if (mActiveServices == null) mActiveServices = new ArrayMap<>();
199            mActiveServices.put(conn, component);
200            conn.onServiceConnected(component, mMockServices.get(component));
201            return true;
202        }
203        return false;
204    }
205
206    @Override
207    public void unbindService(ServiceConnection conn) {
208        if (mService != null) mService.getLeakInfo(conn).clearAllocations();
209        if (mActiveServices != null && mActiveServices.containsKey(conn)) {
210            conn.onServiceDisconnected(mActiveServices.get(conn));
211            mActiveServices.remove(conn);
212            return;
213        }
214        super.unbindService(conn);
215    }
216
217    public boolean isBound(ComponentName component) {
218        return mActiveServices != null && mActiveServices.containsValue(component);
219    }
220
221    @Override
222    public void registerComponentCallbacks(ComponentCallbacks callback) {
223        if (mComponent != null) mComponent.getLeakInfo(callback).addAllocation(new Throwable());
224        super.registerComponentCallbacks(callback);
225    }
226
227    @Override
228    public void unregisterComponentCallbacks(ComponentCallbacks callback) {
229        if (mComponent != null) mComponent.getLeakInfo(callback).clearAllocations();
230        super.unregisterComponentCallbacks(callback);
231    }
232
233    @Override
234    public Statement apply(Statement base, Description description) {
235        return new TestWatcher() {
236            @Override
237            protected void succeeded(Description description) {
238                mSettingsProvider.clearValuesAndCheck(TestableContext.this);
239            }
240
241            @Override
242            protected void failed(Throwable e, Description description) {
243                mSettingsProvider.clearValuesAndCheck(TestableContext.this);
244            }
245        }.apply(base, description);
246    }
247}
248