1 /*
2  * Copyright (C) 2021 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.activitycontext;
18 
19 import static android.content.Intent.FLAG_ACTIVITY_CLEAR_TASK;
20 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
21 
22 import android.app.Activity;
23 import android.app.Instrumentation;
24 import android.content.Context;
25 import android.content.Intent;
26 import android.util.Log;
27 
28 import androidx.test.platform.app.InstrumentationRegistry;
29 
30 import com.android.bedstead.nene.TestApis;
31 import com.android.bedstead.nene.exceptions.NeneException;
32 import com.android.bedstead.nene.users.UserReference;
33 import com.android.compatibility.common.util.ShellIdentityUtils.QuadFunction;
34 import com.android.compatibility.common.util.ShellIdentityUtils.TriFunction;
35 
36 import java.util.concurrent.CountDownLatch;
37 import java.util.concurrent.TimeUnit;
38 import java.util.function.BiConsumer;
39 import java.util.function.BiFunction;
40 import java.util.function.Consumer;
41 import java.util.function.Function;
42 
43 import javax.annotation.Nullable;
44 
45 /**
46  * Activity used for tests which need an actual {@link Context}.
47  */
48 public class ActivityContext extends Activity {
49 
50     private static final String LOG_TAG = "ActivityContext";
51     private static final Context sContext =
52             InstrumentationRegistry.getInstrumentation().getContext();
53 
54     private static Function<Activity, ?> sRunnable;
55     private static @Nullable Object sReturnValue;
56     private static @Nullable Object sThrowValue;
57     private static CountDownLatch sLatch;
58 
59     /**
60      * Run some code using an Activity {@link Context}.
61      *
62      * <p>This method should only be called from an instrumented app.
63      *
64      * <p>The {@link Activity} will be valid within the {@code runnable} callback. Passing the
65      * {@link Activity} outside of the callback is not recommended because it may become invalid
66      * due to lifecycle changes.
67      *
68      * <p>This method will block until the callback has been executed. It will return the same value
69      * as returned by the callback.
70      */
getWithContext(Function<Activity, E> runnable)71     public static <E> E getWithContext(Function<Activity, E> runnable) throws InterruptedException {
72         if (runnable == null) {
73             throw new NullPointerException();
74         }
75 
76         // As we show an Activity we must be in the foreground
77         UserReference currentUser = TestApis.users().current();
78         try {
79             TestApis.users().instrumented().switchTo();
80 
81             Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
82 
83             if (!instrumentation.getContext().getPackageName().equals(
84                     instrumentation.getTargetContext().getPackageName())) {
85                 throw new IllegalStateException(
86                         "ActivityContext can only be used in test apps which instrument themselves."
87                                 + " Consider ActivityScenario for this case.");
88             }
89 
90             synchronized (ActivityContext.class) {
91                 sRunnable = runnable;
92 
93                 sLatch = new CountDownLatch(1);
94                 sReturnValue = null;
95                 sThrowValue = null;
96 
97                 Intent intent = new Intent();
98                 intent.setClass(sContext, ActivityContext.class);
99                 intent.setFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_CLEAR_TASK);
100                 sContext.startActivity(intent);
101             }
102 
103             if (!sLatch.await(5, TimeUnit.MINUTES)) {
104                 throw new NeneException("Timed out while waiting for lambda with context to"
105                         + " complete.");
106             }
107 
108             synchronized (ActivityContext.class) {
109                 sRunnable = null;
110 
111                 if (sThrowValue != null) {
112                     if (sThrowValue instanceof RuntimeException) {
113                         throw (RuntimeException) sThrowValue;
114                     }
115 
116                     if (sThrowValue instanceof Error) {
117                         throw (Error) sThrowValue;
118                     }
119 
120                     throw new IllegalStateException("Invalid value for sThrowValue");
121                 }
122 
123                 return (E) sReturnValue;
124             }
125         } finally {
126             currentUser.switchTo();
127         }
128     }
129 
130     /** {@link #getWithContext(Function)} which does not return a value. */
runWithContext(Consumer<Activity> runnable)131     public static void runWithContext(Consumer<Activity> runnable) throws InterruptedException {
132         getWithContext((inContext) -> {runnable.accept(inContext); return null; });
133     }
134 
135     /** {@link #getWithContext(Function)} with an additional argument. */
getWithContext(E arg1, BiFunction<Activity, E, F> runnable)136     public static <E, F> F getWithContext(E arg1,
137             BiFunction<Activity, E, F> runnable) throws InterruptedException {
138         return getWithContext((inContext) -> runnable.apply(inContext, arg1));
139     }
140 
141     /**
142      * {@link #getWithContext(Function)} which takes an additional argument and does not
143      * return a value.
144      */
runWithContext(E arg1, BiConsumer<Activity, E> runnable)145     public static <E> void runWithContext(E arg1, BiConsumer<Activity, E> runnable)
146             throws InterruptedException {
147         getWithContext((inContext) -> {runnable.accept(inContext, arg1); return null; });
148     }
149 
150     /** {@link #getWithContext(Function)} with two additional arguments. */
getWithContext(E arg1, F arg2, TriFunction<Activity, E, F, G> runnable)151     public static <E, F, G> G getWithContext(E arg1, F arg2,
152             TriFunction<Activity, E, F, G> runnable) throws InterruptedException {
153         return getWithContext((inContext) -> runnable.apply(inContext, arg1, arg2));
154     }
155 
156     /** {@link #getWithContext(Function)} with three additional arguments. */
getWithContext(E arg1, F arg2, G arg3, QuadFunction<Activity, E, F, G, H> runnable)157     public static <E, F, G, H> H getWithContext(E arg1, F arg2, G arg3,
158             QuadFunction<Activity, E, F, G, H> runnable) throws InterruptedException {
159         return getWithContext((inContext) -> runnable.apply(inContext, arg1, arg2, arg3));
160     }
161 
162     @Override
onResume()163     protected void onResume() {
164         super.onResume();
165         synchronized (ActivityContext.class) {
166             if (sRunnable == null) {
167                 Log.e(LOG_TAG, "Launched ActivityContext without runnable");
168             } else {
169                 try {
170                     sReturnValue = sRunnable.apply(this);
171                 } catch (RuntimeException | Error e) {
172                     sThrowValue = e;
173                 }
174                 sLatch.countDown();
175             }
176         }
177     }
178 }
179