From 857d7db69d2aac03ea624a0401e7231a3cef4597 Mon Sep 17 00:00:00 2001
From: Dan Albert <danalbert@google.com>
Date: Tue, 3 Mar 2015 11:47:28 -0800
Subject: [PATCH] Make the root/unroot test more robust.

* Check the current adb user to choose the order of root/unroot.
* Re-root the device when finished.

Change-Id: I47a14b89e2c405bd63722e4d2043fcc629fb5e58
---
 adb/tests/test_adb.py | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/adb/tests/test_adb.py b/adb/tests/test_adb.py
index 49ead732b..f111b0434 100755
--- a/adb/tests/test_adb.py
+++ b/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"