Fixes for RemoteDPC + Tests. am: 5db413da58android-mainline-12.0.0_r99android-mainline-12.0.0_r98android-mainline-12.0.0_r97android-mainline-12.0.0_r96android-mainline-12.0.0_r95android-mainline-12.0.0_r94android-mainline-12.0.0_r93android-mainline-12.0.0_r92android-mainline-12.0.0_r91android-mainline-12.0.0_r90android-mainline-12.0.0_r9android-mainline-12.0.0_r89android-mainline-12.0.0_r88android-mainline-12.0.0_r87android-mainline-12.0.0_r86android-mainline-12.0.0_r85android-mainline-12.0.0_r84android-mainline-12.0.0_r83android-mainline-12.0.0_r82android-mainline-12.0.0_r81android-mainline-12.0.0_r80android-mainline-12.0.0_r8android-mainline-12.0.0_r79android-mainline-12.0.0_r78android-mainline-12.0.0_r77android-mainline-12.0.0_r76android-mainline-12.0.0_r75android-mainline-12.0.0_r74android-mainline-12.0.0_r73android-mainline-12.0.0_r72android-mainline-12.0.0_r71android-mainline-12.0.0_r70android-mainline-12.0.0_r7android-mainline-12.0.0_r69android-mainline-12.0.0_r68android-mainline-12.0.0_r67android-mainline-12.0.0_r66android-mainline-12.0.0_r65android-mainline-12.0.0_r64android-mainline-12.0.0_r63android-mainline-12.0.0_r62android-mainline-12.0.0_r61android-mainline-12.0.0_r60android-mainline-12.0.0_r6android-mainline-12.0.0_r59android-mainline-12.0.0_r58android-mainline-12.0.0_r57android-mainline-12.0.0_r56android-mainline-12.0.0_r53android-mainline-12.0.0_r52android-mainline-12.0.0_r51android-mainline-12.0.0_r50android-mainline-12.0.0_r5android-mainline-12.0.0_r49android-mainline-12.0.0_r48android-mainline-12.0.0_r47android-mainline-12.0.0_r46android-mainline-12.0.0_r45android-mainline-12.0.0_r44android-mainline-12.0.0_r43android-mainline-12.0.0_r42android-mainline-12.0.0_r41android-mainline-12.0.0_r40android-mainline-12.0.0_r39android-mainline-12.0.0_r38android-mainline-12.0.0_r37android-mainline-12.0.0_r35android-mainline-12.0.0_r34android-mainline-12.0.0_r33android-mainline-12.0.0_r32android-mainline-12.0.0_r31android-mainline-12.0.0_r30android-mainline-12.0.0_r3android-mainline-12.0.0_r29android-mainline-12.0.0_r28android-mainline-12.0.0_r27android-mainline-12.0.0_r26android-mainline-12.0.0_r25android-mainline-12.0.0_r24android-mainline-12.0.0_r23android-mainline-12.0.0_r22android-mainline-12.0.0_r21android-mainline-12.0.0_r20android-mainline-12.0.0_r2android-mainline-12.0.0_r19android-mainline-12.0.0_r18android-mainline-12.0.0_r17android-mainline-12.0.0_r16android-mainline-12.0.0_r15android-mainline-12.0.0_r14android-mainline-12.0.0_r13android-mainline-12.0.0_r126android-mainline-12.0.0_r125android-mainline-12.0.0_r124android-mainline-12.0.0_r123android-mainline-12.0.0_r122android-mainline-12.0.0_r121android-mainline-12.0.0_r120android-mainline-12.0.0_r12android-mainline-12.0.0_r119android-mainline-12.0.0_r118android-mainline-12.0.0_r117android-mainline-12.0.0_r116android-mainline-12.0.0_r115android-mainline-12.0.0_r114android-mainline-12.0.0_r113android-mainline-12.0.0_r110android-mainline-12.0.0_r11android-mainline-12.0.0_r109android-mainline-12.0.0_r108android-mainline-12.0.0_r107android-mainline-12.0.0_r106android-mainline-12.0.0_r105android-mainline-12.0.0_r104android-mainline-12.0.0_r103android-mainline-12.0.0_r102android-mainline-12.0.0_r101android-mainline-12.0.0_r100android-mainline-12.0.0_r10android-mainline-12.0.0_r1aml_wif_311811030aml_tz3_311312010aml_tet_311811050aml_sdk_311710000aml_pco_311011000aml_mpr_311911090aml_doc_310851020android12-mainline-wifi-releaseandroid12-mainline-tethering-releaseandroid12-mainline-statsd-releaseandroid12-mainline-sdkext-releaseandroid12-mainline-resolv-releaseandroid12-mainline-permission-releaseandroid12-mainline-neuralnetworks-releaseandroid12-mainline-networkstack-releaseandroid12-mainline-mediaprovider-releaseandroid12-mainline-media-swcodec-releaseandroid12-mainline-media-releaseandroid12-mainline-ipsec-releaseandroid12-mainline-extservices-releaseandroid12-mainline-documentsui-releaseandroid12-mainline-conscrypt-releaseandroid12-mainline-cellbroadcast-releaseandroid12-mainline-captiveportallogin-releaseandroid12-mainline-art-releaseandroid12-mainline-adbd-release

Original change: https://googleplex-android-review.googlesource.com/c/platform/external/connectedappssdk/+/14639253

Change-Id: Ib8a3bb4ee3a9f9dfb690a8e3aff35a3e23b0a172

9 files changed, 166 insertions, 105 deletions

diff --git a/processor/src/main/java/com/google/android/enterprise/connectedapps/processor/DispatcherGenerator.java b/processor/src/main/java/com/google/android/enterprise/connectedapps/processor/DispatcherGenerator.java
index ff45251..f5264b0 100644
--- a/processor/src/main/java/com/google/android/enterprise/connectedapps/processor/DispatcherGenerator.java
+++ b/processor/src/main/java/com/google/android/enterprise/connectedapps/processor/DispatcherGenerator.java
@@ -92,7 +92,6 @@ final class DispatcherGenerator {
             .initializer("new $T()", PARCEL_CALL_RECEIVER_CLASSNAME)
             .build());
 
-    addEnsureValidCallerMethod(classBuilder);
     addCallMethod(classBuilder);
     addPrepareCallMethod(classBuilder);
     addFetchResponseMethod(classBuilder);
@@ -109,7 +108,6 @@ final class DispatcherGenerator {
             .addParameter(int.class, "blockId")
             .addParameter(int.class, "numBytes")
             .addParameter(ArrayTypeName.of(byte.class), "paramBytes")
-            .addStatement("ensureValidCaller(context)")
             .addStatement("parcelCallReceiver.prepareCall(callId, blockId, numBytes, paramBytes)")
             .addJavadoc(
                 "Store a block of bytes to be part of a future call to\n"
@@ -142,7 +140,6 @@ final class DispatcherGenerator {
             .addParameter(long.class, "callId")
             .addParameter(int.class, "blockId")
             .returns(ArrayTypeName.of(byte.class))
-            .addStatement("ensureValidCaller(context)")
             .addStatement("return parcelCallReceiver.getPreparedResponse(callId, blockId)")
             .addJavadoc(
                 "Fetch a response block if a previous call to\n {@link #call(Context, long, int,"
@@ -157,41 +154,11 @@ final class DispatcherGenerator {
     classBuilder.addMethod(prepareCallMethod);
   }
 
-  private static void addEnsureValidCallerMethod(TypeSpec.Builder classBuilder) {
-    CodeBlock.Builder methodCode = CodeBlock.builder();
-
-    methodCode.addStatement(
-        "$T[] callingPackageNames ="
-            + " context.getPackageManager().getPackagesForUid($T.getCallingUid())",
-        String.class,
-        BINDER_CLASSNAME);
-    methodCode.beginControlFlow("for (String callingPackageName : callingPackageNames)");
-    methodCode.beginControlFlow("if (context.getPackageName().equals(callingPackageName))");
-    methodCode.addStatement("return");
-    methodCode.endControlFlow();
-    methodCode.endControlFlow();
-
-    methodCode.addStatement(
-        "throw new $T(\"Cross-profile functionality is only available within the same package\")",
-        IllegalStateException.class);
-
-    MethodSpec ensureValidCallerMethod =
-        MethodSpec.methodBuilder("ensureValidCaller")
-            .addModifiers(Modifier.PRIVATE, Modifier.STATIC)
-            .addParameter(CONTEXT_CLASSNAME, "context")
-            .addCode(methodCode.build())
-            .build();
-
-    classBuilder.addMethod(ensureValidCallerMethod);
-  }
-
   private void addCallMethod(TypeSpec.Builder classBuilder) {
     CodeBlock.Builder methodCode = CodeBlock.builder();
 
     methodCode.beginControlFlow("try");
 
-    methodCode.addStatement("ensureValidCaller(context)");
-
     methodCode.addStatement(
         "$1T parcel = parcelCallReceiver.getPreparedCall(callId, blockId, paramBytes)",
         PARCEL_CLASSNAME);
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/BothProfilesTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/BothProfilesTest.java
index b90ae24..8b15d35 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/BothProfilesTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/BothProfilesTest.java
@@ -26,6 +26,9 @@ import com.google.android.enterprise.connectedapps.testapp.connector.TestProfile
 import com.google.android.enterprise.connectedapps.testapp.types.ProfileTestCrossProfileTypeWhichNeedsContext;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
+
+import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -39,8 +42,8 @@ public class BothProfilesTest {
 
   private static final Application context = ApplicationProvider.getApplicationContext();
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
-  private final InstrumentedTestUtilities utilities =
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
 
   private final ProfileTestCrossProfileTypeWhichNeedsContext type =
@@ -51,6 +54,11 @@ public class BothProfilesTest {
     utilities.ensureReadyForCrossProfileCalls();
   }
 
+  @AfterClass
+  public static void teardown() {
+    utilities.ensureNoWorkProfile();
+  }
+
   /** This test could not be covered by Robolectric. */
   @Test
   public void both_synchronous_timesOutOnWorkProfile_timeoutNotEnforcedOnSynchronousCalls() {
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/ConnectTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/ConnectTest.java
index 0ab9c44..84069a7 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/ConnectTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/ConnectTest.java
@@ -23,6 +23,9 @@ import androidx.test.core.app.ApplicationProvider;
 import com.google.android.enterprise.connectedapps.exceptions.UnavailableProfileException;
 import com.google.android.enterprise.connectedapps.instrumented.utils.InstrumentedTestUtilities;
 import com.google.android.enterprise.connectedapps.testapp.connector.TestProfileConnector;
+
+import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,8 +43,8 @@ import org.junit.runners.JUnit4;
 public class ConnectTest {
   private static final Application context = ApplicationProvider.getApplicationContext();
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
-  private final InstrumentedTestUtilities utilities =
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
 
   @Before
@@ -49,6 +52,17 @@ public class ConnectTest {
     utilities.ensureReadyForCrossProfileCalls();
   }
 
+  @After
+  public void teardown() {
+    connector.stopManualConnectionManagement();
+    utilities.waitForDisconnected();
+  }
+
+  @AfterClass
+  public static void teardownClass() {
+    utilities.ensureNoWorkProfile();
+  }
+
   @Test
   public void connect_connects() throws Exception {
     utilities.ensureReadyForCrossProfileCalls();
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/HappyPathEndToEndTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/HappyPathEndToEndTest.java
index de542d1..3b9ea57 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/HappyPathEndToEndTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/HappyPathEndToEndTest.java
@@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import android.app.Application;
 import androidx.test.core.app.ApplicationProvider;
+
 import com.google.android.enterprise.connectedapps.exceptions.UnavailableProfileException;
 import com.google.android.enterprise.connectedapps.instrumented.utils.BlockingExceptionCallbackListener;
 import com.google.android.enterprise.connectedapps.instrumented.utils.BlockingStringCallbackListener;
@@ -29,10 +30,13 @@ import com.google.android.enterprise.connectedapps.testapp.types.ProfileTestCros
 import com.google.android.enterprise.connectedapps.testing.BlockingPoll;
 import java.util.concurrent.ExecutionException;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import android.util.Log;
 
 /**
  * Tests for high level behaviour running on a correctly configured device (with a managed profile
@@ -46,8 +50,8 @@ public class HappyPathEndToEndTest {
 
   private static final String STRING = "String";
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
-  private final InstrumentedTestUtilities utilities =
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
   private final ProfileTestCrossProfileType type = ProfileTestCrossProfileType.create(connector);
   private final ProfileTestCrossProfileTypeWhichNeedsContext typeWithContext =
@@ -64,6 +68,11 @@ public class HappyPathEndToEndTest {
     utilities.waitForDisconnected();
   }
 
+  @AfterClass
+  public static void teardownClass() {
+    utilities.ensureNoWorkProfile();
+  }
+
   @Test
   public void isAvailable_isTrue() {
     assertThat(connector.isAvailable()).isTrue();
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/InstrumentedTestUtilitiesTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/InstrumentedTestUtilitiesTest.java
index 7a84852..81c9ce8 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/InstrumentedTestUtilitiesTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/InstrumentedTestUtilitiesTest.java
@@ -22,6 +22,8 @@ import androidx.test.core.app.ApplicationProvider;
 import com.google.android.enterprise.connectedapps.AvailabilityListener;
 import com.google.android.enterprise.connectedapps.testapp.connector.TestProfileConnector;
 import com.google.android.enterprise.connectedapps.testing.InstrumentedTestUtilities;
+
+import org.junit.AfterClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -32,10 +34,15 @@ public class InstrumentedTestUtilitiesTest {
 
   private static final Application context = ApplicationProvider.getApplicationContext();
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
-  private final InstrumentedTestUtilities utilities =
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
 
+  @AfterClass
+  public static void teardown() {
+    utilities.ensureNoWorkProfile();
+  }
+
   @Test
   public void isAvailable_ensureReadyForCrossProfileCalls_isTrue() {
     utilities.ensureReadyForCrossProfileCalls();
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/MessageSizeTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/MessageSizeTest.java
index 25be188..363b271 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/MessageSizeTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/MessageSizeTest.java
@@ -27,6 +27,8 @@ import com.google.android.enterprise.connectedapps.instrumented.utils.Instrument
 import com.google.android.enterprise.connectedapps.testapp.connector.TestProfileConnector;
 import com.google.android.enterprise.connectedapps.testapp.types.ProfileTestCrossProfileType;
 import java.util.concurrent.ExecutionException;
+
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,8 +42,8 @@ public class MessageSizeTest {
   private static final String SMALL_STRING = "String";
   private static final String LARGE_STRING = randomString(1500000); // 3Mb
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
-  private final InstrumentedTestUtilities utilities =
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
   private final ProfileTestCrossProfileType type = ProfileTestCrossProfileType.create(connector);
 
@@ -55,6 +57,11 @@ public class MessageSizeTest {
     utilities.ensureReadyForCrossProfileCalls();
   }
 
+  @AfterClass
+  public static void teardown() {
+    utilities.ensureNoWorkProfile();
+  }
+
   @Test
   public void synchronous_smallMessage_sends() throws UnavailableProfileException {
     utilities.manuallyConnectAndWait();
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotInstalledInOtherUserTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotInstalledInOtherUserTest.java
index b3c4b38..9b6de41 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotInstalledInOtherUserTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotInstalledInOtherUserTest.java
@@ -25,6 +25,8 @@ import com.google.android.enterprise.connectedapps.instrumented.utils.Instrument
 import com.google.android.enterprise.connectedapps.testapp.connector.TestProfileConnector;
 import com.google.android.enterprise.connectedapps.testapp.types.ProfileTestCrossProfileType;
 import com.google.common.util.concurrent.ListenableFuture;
+
+import org.junit.AfterClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -40,11 +42,16 @@ public class NotInstalledInOtherUserTest {
 
   private static final Application context = ApplicationProvider.getApplicationContext();
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
   private final ProfileTestCrossProfileType type = ProfileTestCrossProfileType.create(connector);
-  private final InstrumentedTestUtilities utilities =
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
 
+  @AfterClass
+  public static void teardown() {
+    utilities.ensureNoWorkProfile();
+  }
+
   @Test
   public void asyncCall_notInstalledInOtherProfile_failsFast() {
     utilities.ensureWorkProfileExistsWithoutTestApp();
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotReallySerializableTest.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotReallySerializableTest.java
index abb5133..912138b 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotReallySerializableTest.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/tests/NotReallySerializableTest.java
@@ -29,6 +29,8 @@ import com.google.android.enterprise.connectedapps.testapp.NotReallySerializable
 import com.google.android.enterprise.connectedapps.testapp.connector.TestProfileConnector;
 import com.google.android.enterprise.connectedapps.testapp.types.ProfileTestCrossProfileType;
 import com.google.common.util.concurrent.ListenableFuture;
+
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -44,8 +46,8 @@ import org.junit.runners.JUnit4;
 public class NotReallySerializableTest {
   private static final Application context = ApplicationProvider.getApplicationContext();
 
-  private final TestProfileConnector connector = TestProfileConnector.create(context);
-  private final InstrumentedTestUtilities utilities =
+  private static final TestProfileConnector connector = TestProfileConnector.create(context);
+  private static final InstrumentedTestUtilities utilities =
       new InstrumentedTestUtilities(context, connector);
 
   private final ProfileTestCrossProfileType type = ProfileTestCrossProfileType.create(connector);
@@ -57,6 +59,11 @@ public class NotReallySerializableTest {
     utilities.ensureReadyForCrossProfileCalls();
   }
 
+  @AfterClass
+  public static void teardown() {
+    utilities.ensureNoWorkProfile();
+  }
+
   @Test
   public void
       synchronous_serializableObjectIsNotReallySerializable_throwsProfileRuntimeException() {
diff --git a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/utils/InstrumentedTestUtilities.java b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/utils/InstrumentedTestUtilities.java
index e76ee9c..be1db8e 100644
--- a/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/utils/InstrumentedTestUtilities.java
+++ b/tests/instrumented/src/main/java/com/google/android/enterprise/connectedapps/instrumented/utils/InstrumentedTestUtilities.java
@@ -44,54 +44,73 @@ public class InstrumentedTestUtilities {
   private final ProfileConnector connector;
   private final Context context;
   private final com.google.android.enterprise.connectedapps.testing.InstrumentedTestUtilities
-      instrumentedTestUtilities;
+          instrumentedTestUtilities;
 
+  private static final int S_REQUEST_QUIET_MODE_ENABLED_ID = 73;
   private static final int R_REQUEST_QUIET_MODE_ENABLED_ID = 72;
   private static final int REQUEST_QUIET_MODE_ENABLED_ID = 58;
 
   private static final String USER_ID_KEY = "USER_ID";
   private static final Parameter USER_ID_PARAMETER = new Parameter(USER_ID_KEY);
 
+  private static final ServiceCall S_TURN_OFF_WORK_PROFILE_COMMAND =
+          new ServiceCall("user", S_REQUEST_QUIET_MODE_ENABLED_ID)
+                  .setUser(1000) // user 1000 has packageName "android"
+                  .addStringParam("android") // callingPackage
+                  .addBooleanParam(true) // enableQuietMode
+                  .addIntParam(USER_ID_PARAMETER) // userId
+                  .addIntParam(0) // target
+                  .addIntParam(0); // flags
+
   private static final ServiceCall R_TURN_OFF_WORK_PROFILE_COMMAND =
-      new ServiceCall("user", R_REQUEST_QUIET_MODE_ENABLED_ID)
-          .setUser(1000) // user 1000 has packageName "android"
-          .addStringParam("android") // callingPackage
-          .addBooleanParam(true) // enableQuietMode
-          .addIntParam(USER_ID_PARAMETER) // userId
-          .addIntParam(0) // target
-          .addIntParam(0); // flags
+          new ServiceCall("user", R_REQUEST_QUIET_MODE_ENABLED_ID)
+                  .setUser(1000) // user 1000 has packageName "android"
+                  .addStringParam("android") // callingPackage
+                  .addBooleanParam(true) // enableQuietMode
+                  .addIntParam(USER_ID_PARAMETER) // userId
+                  .addIntParam(0) // target
+                  .addIntParam(0); // flags
 
   private static final ServiceCall TURN_OFF_WORK_PROFILE_COMMAND =
-      new ServiceCall("user", REQUEST_QUIET_MODE_ENABLED_ID)
-          .setUser(1000) // user 1000 has packageName "android"
-          .addStringParam("android") // callingPackage
-          .addBooleanParam(true) // enableQuietMode
-          .addIntParam(USER_ID_PARAMETER) // userId
-          .addIntParam(0); // target
+          new ServiceCall("user", REQUEST_QUIET_MODE_ENABLED_ID)
+                  .setUser(1000) // user 1000 has packageName "android"
+                  .addStringParam("android") // callingPackage
+                  .addBooleanParam(true) // enableQuietMode
+                  .addIntParam(USER_ID_PARAMETER) // userId
+                  .addIntParam(0); // target
+
+  private static final ServiceCall S_TURN_ON_WORK_PROFILE_COMMAND =
+          new ServiceCall("user", S_REQUEST_QUIET_MODE_ENABLED_ID)
+                  .setUser(1000) // user 1000 has packageName "android"
+                  .addStringParam("android") // callingPackage
+                  .addBooleanParam(false) // enableQuietMode
+                  .addIntParam(USER_ID_PARAMETER) // userId
+                  .addIntParam(0) // target
+                  .addIntParam(0); // flags
 
   private static final ServiceCall R_TURN_ON_WORK_PROFILE_COMMAND =
-      new ServiceCall("user", R_REQUEST_QUIET_MODE_ENABLED_ID)
-          .setUser(1000) // user 1000 has packageName "android"
-          .addStringParam("android") // callingPackage
-          .addBooleanParam(false) // enableQuietMode
-          .addIntParam(USER_ID_PARAMETER) // userId
-          .addIntParam(0) // target
-          .addIntParam(0); // flags
+          new ServiceCall("user", R_REQUEST_QUIET_MODE_ENABLED_ID)
+                  .setUser(1000) // user 1000 has packageName "android"
+                  .addStringParam("android") // callingPackage
+                  .addBooleanParam(false) // enableQuietMode
+                  .addIntParam(USER_ID_PARAMETER) // userId
+                  .addIntParam(0) // target
+                  .addIntParam(0); // flags
 
   private static final ServiceCall TURN_ON_WORK_PROFILE_COMMAND =
-      new ServiceCall("user", REQUEST_QUIET_MODE_ENABLED_ID)
-          .setUser(1000) // user 1000 has packageName "android"
-          .addStringParam("android") // callingPackage
-          .addBooleanParam(false) // enableQuietMode
-          .addIntParam(USER_ID_PARAMETER) // userId
-          .addIntParam(0); // target
+          new ServiceCall("user", REQUEST_QUIET_MODE_ENABLED_ID)
+                  .setUser(1000) // user 1000 has packageName "android"
+                  .addStringParam("android") // callingPackage
+                  .addBooleanParam(false) // enableQuietMode
+                  .addIntParam(USER_ID_PARAMETER) // userId
+                  .addIntParam(0); // target
 
   public InstrumentedTestUtilities(Context context, ProfileConnector connector) {
     this.context = context;
     this.connector = connector;
     this.instrumentedTestUtilities =
-        new com.google.android.enterprise.connectedapps.testing.InstrumentedTestUtilities(
-            context, connector);
+            new com.google.android.enterprise.connectedapps.testing.InstrumentedTestUtilities(
+                    context, connector);
   }
 
   /**
@@ -137,7 +156,7 @@ public class InstrumentedTestUtilities {
 
   public void installInUser(int userId) {
     runCommandWithOutput(
-        "cmd package install-existing --user " + userId + " " + context.getPackageName());
+            "cmd package install-existing --user " + userId + " " + context.getPackageName());
   }
 
   /**
@@ -148,7 +167,7 @@ public class InstrumentedTestUtilities {
   public void grantInteractAcrossUsers() {
     // TODO(scottjonathan): Support INTERACT_ACROSS_PROFILES in these tests.
     runCommandWithOutput(
-        "pm grant " + context.getPackageName() + " android.permission.INTERACT_ACROSS_USERS");
+            "pm grant " + context.getPackageName() + " android.permission.INTERACT_ACROSS_USERS");
   }
 
   /**
@@ -171,12 +190,15 @@ public class InstrumentedTestUtilities {
         return;
       }
 
+
+      runCommandWithOutput("pm remove-user " + getWorkProfileUserId());
+
       // TODO(162219825): Try to remove the package
 
-      throw new IllegalStateException(
-          "There is already a work profile on the device with user id "
-              + getWorkProfileUserId()
-              + ".");
+//      throw new IllegalStateException(
+//              "There is already a work profile on the device with user id "
+//                      + getWorkProfileUserId()
+//                      + ".");
     }
     runCommandWithOutput("pm create-user --profileOf 0 --managed TestProfile123");
     int workProfileUserId = getWorkProfileUserId();
@@ -186,7 +208,7 @@ public class InstrumentedTestUtilities {
   private static boolean userHasPackageInstalled(int userId, String packageName) {
     String expectedPackageLine = "package:" + packageName;
     String[] installedPackages =
-        runCommandWithOutput("pm list packages --user " + userId).split("\n");
+            runCommandWithOutput("pm list packages --user " + userId).split("\n");
     for (String installedPackage : installedPackages) {
       if (installedPackage.equals(expectedPackageLine)) {
         return true;
@@ -227,18 +249,24 @@ public class InstrumentedTestUtilities {
    * @see #turnOffWorkProfileAndWait()
    */
   public void turnOffWorkProfile() {
-    if (VERSION.SDK_INT == VERSION_CODES.R) {
+    if (VERSION.CODENAME.equals("S")) {
       runCommandWithOutput(
-          R_TURN_OFF_WORK_PROFILE_COMMAND
-              .prepare()
-              .setInt(USER_ID_KEY, getWorkProfileUserId())
-              .getCommand());
+              S_TURN_OFF_WORK_PROFILE_COMMAND
+                      .prepare()
+                      .setInt(USER_ID_KEY, getWorkProfileUserId())
+                      .getCommand());
+    } else if (VERSION.SDK_INT == VERSION_CODES.R) {
+      runCommandWithOutput(
+              R_TURN_OFF_WORK_PROFILE_COMMAND
+                      .prepare()
+                      .setInt(USER_ID_KEY, getWorkProfileUserId())
+                      .getCommand());
     } else if (VERSION.SDK_INT == VERSION_CODES.Q || VERSION.SDK_INT == VERSION_CODES.P) {
       runCommandWithOutput(
-          TURN_OFF_WORK_PROFILE_COMMAND
-              .prepare()
-              .setInt(USER_ID_KEY, getWorkProfileUserId())
-              .getCommand());
+              TURN_OFF_WORK_PROFILE_COMMAND
+                      .prepare()
+                      .setInt(USER_ID_KEY, getWorkProfileUserId())
+                      .getCommand());
     } else {
       throw new IllegalStateException("Cannot turn off work on this version of android");
     }
@@ -259,7 +287,7 @@ public class InstrumentedTestUtilities {
     turnOnWorkProfile();
 
     ProfileAvailabilityPoll.blockUntilProfileRunningAndUnlocked(
-        context, getWorkProfileUserHandle());
+            context, getWorkProfileUserHandle());
   }
 
   // TODO(160147511): Remove use of service calls for versions after R
@@ -271,18 +299,24 @@ public class InstrumentedTestUtilities {
    * @see #turnOnWorkProfileAndWait()
    */
   public void turnOnWorkProfile() {
-    if (VERSION.SDK_INT == VERSION_CODES.R) {
+    if (VERSION.CODENAME.equals("S")) {
+      runCommandWithOutput(
+              S_TURN_ON_WORK_PROFILE_COMMAND
+                      .prepare()
+                      .setInt(USER_ID_KEY, getWorkProfileUserId())
+                      .getCommand());
+    } else if (VERSION.SDK_INT == VERSION_CODES.R) {
       runCommandWithOutput(
-          R_TURN_ON_WORK_PROFILE_COMMAND
-              .prepare()
-              .setInt(USER_ID_KEY, getWorkProfileUserId())
-              .getCommand());
+              R_TURN_ON_WORK_PROFILE_COMMAND
+                      .prepare()
+                      .setInt(USER_ID_KEY, getWorkProfileUserId())
+                      .getCommand());
     } else if (VERSION.SDK_INT == VERSION_CODES.Q || VERSION.SDK_INT == VERSION_CODES.P) {
       runCommandWithOutput(
-          TURN_ON_WORK_PROFILE_COMMAND
-              .prepare()
-              .setInt(USER_ID_KEY, getWorkProfileUserId())
-              .getCommand());
+              TURN_ON_WORK_PROFILE_COMMAND
+                      .prepare()
+                      .setInt(USER_ID_KEY, getWorkProfileUserId())
+                      .getCommand());
     } else {
       throw new IllegalStateException("Cannot turn on work on this version of android");
     }
@@ -314,7 +348,7 @@ public class InstrumentedTestUtilities {
   }
 
   private static final Pattern CREATE_USER_PATTERN =
-      Pattern.compile("Success: created user id (\\d+)");
+          Pattern.compile("Success: created user id (\\d+)");
 
   public int createUser(String username) {
     String output = runCommandWithOutput("pm create-user " + username);
@@ -334,6 +368,7 @@ public class InstrumentedTestUtilities {
   }
 
   private static String runCommandWithOutput(String command) {
+    // TODO: Log output so we can see why it's failing
     ParcelFileDescriptor p = runCommand(command);
 
     InputStream inputStream = new FileInputStream(p.getFileDescriptor());
@@ -347,7 +382,7 @@ public class InstrumentedTestUtilities {
 
   private static ParcelFileDescriptor runCommand(String command) {
     return InstrumentationRegistry.getInstrumentation()
-        .getUiAutomation()
-        .executeShellCommand(command);
+            .getUiAutomation()
+            .executeShellCommand(command);
   }
 }