Merge "Make the root/unroot test more robust."

adb/tests/test_adb.py

@@ -237,16 +237,36 @@ class AdbBasic(unittest.TestCase):
                 version_num = True
         self.assertTrue(version_num)
 
-    def test_root_unroot(self):
-        """Make sure that adb root and adb unroot work, using id(1)."""
+    def _test_root(self):
         adb = AdbWrapper()
         adb.root()
         adb.wait()
         self.assertEqual("root", adb.shell("id -un").strip())
+
+    def _test_unroot(self):
+        adb = AdbWrapper()
         adb.unroot()
         adb.wait()
         self.assertEqual("shell", adb.shell("id -un").strip())
 
+    def test_root_unroot(self):
+        """Make sure that adb root and adb unroot work, using id(1)."""
+        adb = AdbWrapper()
+        original_user = adb.shell("id -un").strip()
+        try:
+            if original_user == "root":
+                self._test_unroot()
+                self._test_root()
+            elif original_user == "shell":
+                self._test_root()
+                self._test_unroot()
+        finally:
+            if original_user == "root":
+                adb.root()
+            else:
+                adb.unroot()
+            adb.wait()
+
 
 class AdbFile(unittest.TestCase):
     SCRATCH_DIR = "/data/local/tmp"