Merge "Add new class for SetPropRule and move outside ShellCommandRule." into main
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java b/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
index 79724e1..19c5850 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
@@ -15,8 +15,6 @@
  */
 package com.android.launcher3.ui.widget;
 
-import static com.android.launcher3.util.rule.ShellCommandRule.createEnableInputTransportPublisherRule;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
@@ -32,9 +30,7 @@
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.launcher3.widget.picker.WidgetsRecyclerView;
 
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 
 /**
@@ -44,9 +40,6 @@
 @MediumTest
 @RunWith(AndroidJUnit4.class)
 public class TaplWidgetPickerTest extends AbstractLauncherUiTest<Launcher> {
-    // b/325377690 : To get the log printed where DOWN key event is getting lost from TAPL.
-    @Rule public final TestRule mEnableInputTransportPublisherRule =
-            createEnableInputTransportPublisherRule();
 
     private WidgetsRecyclerView getWidgetsView(Launcher launcher) {
         return WidgetsFullSheet.getWidgetsView(launcher);
diff --git a/tests/src/com/android/launcher3/util/rule/SetPropRule.java b/tests/src/com/android/launcher3/util/rule/SetPropRule.java
new file mode 100644
index 0000000..74fec35
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/SetPropRule.java
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util.rule;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import android.text.TextUtils;
+
+import androidx.annotation.NonNull;
+import androidx.test.uiautomator.UiDevice;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+/**
+ * Test rule which executes a set prop command at the start of the test.
+ * This rule needs the property tag and property value so that value can be set to a tag.
+ */
+public class SetPropRule implements TestRule {
+    private static final String SETPROP_PREFIX = "setprop";
+    private static final String GETPROP_PREFIX = "getprop";
+    private static final String UNKNOWN = "UNKNOWN";
+    @NonNull private final String mPropTag;
+    @NonNull private final String mPropValue;
+
+    public SetPropRule(@NonNull String propTag, @NonNull String propValue) {
+        mPropTag = propTag.trim();
+        mPropValue = propValue.trim();
+    }
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        return new Statement() {
+            @Override
+            public void evaluate() throws Throwable {
+                String getpropCmd = GETPROP_PREFIX + " " + mPropTag;
+                String initialValue = UiDevice.getInstance(getInstrumentation())
+                        .executeShellCommand(getpropCmd);
+                if (TextUtils.isEmpty(initialValue.trim())) {
+                    initialValue = UNKNOWN;
+                }
+                // setprop command always follows format : setprop <TAG> <value>
+                String revertSetPropCmd = SETPROP_PREFIX + " " + mPropTag + " " + initialValue;
+                String setPropCmd = SETPROP_PREFIX + " " + mPropTag + " " + mPropValue;
+                new ShellCommandRule(setPropCmd, revertSetPropCmd)
+                        .apply(base, description).evaluate();
+            }
+        };
+    }
+
+    /**
+     * Enables "InputTransportPublisher" debug flag. This prints the key input events dispatched by
+     * the system server.
+     * adb shell setprop log.tag.InputTransportPublisher DEBUG
+     * See {@link com.android.cts.input.DebugInputRule} for more details.
+     */
+    public static SetPropRule createEnableInputTransportPublisherRule() {
+        return new SetPropRule("log.tag.InputTransportPublisher", "DEBUG");
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
index 977995e..2219003 100644
--- a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
@@ -15,13 +15,12 @@
  */
 package com.android.launcher3.util.rule;
 
-import static androidx.test.InstrumentationRegistry.getInstrumentation;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
 
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
-import android.text.TextUtils;
 
 import androidx.annotation.Nullable;
 import androidx.test.InstrumentationRegistry;
@@ -40,10 +39,6 @@
  * Test rule which executes a shell command at the start of the test.
  */
 public class ShellCommandRule implements TestRule {
-
-    private static final String SETPROP_PREFIX = "setprop";
-    private static final String GETPROP_PREFIX = "getprop";
-    private static final String UNKNOWN = "UNKNOWN";
     private final String mCmd;
     private final String mRevertCommand;
     private final boolean mCheckSuccess;
@@ -66,19 +61,6 @@
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
-                String revertSetPropCmd = null;
-                if (mCmd.startsWith(SETPROP_PREFIX) && mRevertCommand == null) {
-                    // setprop command always follows format : setprop <TAG> <value>
-                    // We are stripping out only the TAG here
-                    String tag = mCmd.split("\\s")[1];
-                    String getpropCmd = GETPROP_PREFIX + " " + tag;
-                    String initialValue = UiDevice.getInstance(
-                            getInstrumentation()).executeShellCommand(getpropCmd);
-                    if (TextUtils.isEmpty(initialValue.trim())) {
-                        initialValue = UNKNOWN;
-                    }
-                    revertSetPropCmd = SETPROP_PREFIX + " " + tag + " " + initialValue;
-                }
                 final String result =
                         UiDevice.getInstance(getInstrumentation()).executeShellCommand(mCmd);
                 if (mCheckSuccess) {
@@ -90,15 +72,12 @@
                 try {
                     base.evaluate();
                 } finally {
-                    if (mRevertCommand != null || revertSetPropCmd != null) {
-                        String revertCmd =
-                                mRevertCommand != null ? mRevertCommand : revertSetPropCmd;
+                    if (mRevertCommand != null) {
                         final String revertResult = UiDevice.getInstance(
-                                getInstrumentation()).executeShellCommand(
-                                revertCmd);
+                                getInstrumentation()).executeShellCommand(mRevertCommand);
                         if (mCheckSuccess) {
                             Assert.assertTrue(
-                                    "Failed command: " + revertCmd
+                                    "Failed command: " + mRevertCommand
                                             + ", result: " + revertResult,
                                     "Success".equals(result.replaceAll("\\s", "")));
                         }
@@ -141,14 +120,4 @@
         return new ShellCommandRule("settings put global heads_up_notifications_enabled 0",
                 "settings put global heads_up_notifications_enabled 1");
     }
-
-    /**
-     * Enables "InputTransportPublisher" debug flag. This prints the key input events dispatched by
-     * the system server.
-     * adb shell setprop log.tag.InputTransportPublisher DEBUG
-     * See {@link com.android.cts.input.DebugInputRule} for more details.
-     */
-    public static ShellCommandRule createEnableInputTransportPublisherRule() {
-        return new ShellCommandRule("setprop log.tag.InputTransportPublisher DEBUG", null);
-    }
 }