platform_CompromisedStatefulPartition: clear the TPM before test ends

This change clears the TPM at the end of the test, even if an error is
raised.

Before this change, the system would be in a weird state that
tpm_manager db is gone whereas the TPM remains owned. Not sure if
other daemons/apps are affected.

This change might solve the empty lockout password problem in tests
that soft-clear the TPM. See b/148752436 for details.

BUG=b:148752436
TEST=test_that -b soraka <DUP IP> platform_CompromisedStatefulPartition

Change-Id: Ic73ffeef12de7446cc3554ee5963ad532e5212d9
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2046683
Tested-by: Wei-Cheng Xiao <garryxiao@chromium.org>
Reviewed-by: Andrey Pronin <apronin@chromium.org>
Commit-Queue: Wei-Cheng Xiao <garryxiao@chromium.org>
diff --git a/server/site_tests/platform_CompromisedStatefulPartition/control b/server/site_tests/platform_CompromisedStatefulPartition/control
index 7e8dd75..e2684aa 100644
--- a/server/site_tests/platform_CompromisedStatefulPartition/control
+++ b/server/site_tests/platform_CompromisedStatefulPartition/control
@@ -30,6 +30,6 @@
 def run(machine):
     host = hosts.create_host(machine)
     job.run_test("platform_CompromisedStatefulPartition", host=host,
-                 disable_sysinfo=True, client_autotest="desktopui_SimpleLogin")
+                 disable_sysinfo=True, client_test="desktopui_SimpleLogin")
 
 parallel_simple(run, machines)
diff --git a/server/site_tests/platform_CompromisedStatefulPartition/platform_CompromisedStatefulPartition.py b/server/site_tests/platform_CompromisedStatefulPartition/platform_CompromisedStatefulPartition.py
index faddb5e..399e71e 100644
--- a/server/site_tests/platform_CompromisedStatefulPartition/platform_CompromisedStatefulPartition.py
+++ b/server/site_tests/platform_CompromisedStatefulPartition/platform_CompromisedStatefulPartition.py
@@ -2,10 +2,23 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from autotest_lib.client.common_lib import error
-from autotest_lib.server import autotest, test
+import contextlib
 import time
 
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib.cros import tpm_utils
+from autotest_lib.server import autotest, test
+
+@contextlib.contextmanager
+def ensure_tpm_reset(client):
+    """This context manager ensures we clears the TPM and restore the system to
+    a reasonable state at the end of the test, even when an error occurs.
+    """
+    try:
+        yield
+    finally:
+        tpm_utils.ClearTPMOwnerRequest(client)
+
 class platform_CompromisedStatefulPartition(test.test):
     """Tests how the system recovers with the corrupted stateful partition.
     """
@@ -24,22 +37,13 @@
         '/mnt/stateful_partition/encrypted.needs-finalization',
     ]
 
-
-    def run_once(self, host, client_autotest):
-        """This test verify that user should get OOBE after booting
-        the device with corrupted stateful partition.
-        Test fails if not able to recover the device with corrupted
-        stateful partition.
+    def _test_stateful_corruption(self, autotest_client, host, client_test):
+        """Corrupts the stateful partition and reboots; checks if the client
+        goes through OOBE again and if encstateful is recreated.
         """
-        if host.get_board_type() == 'OTHER':
-            raise error.TestNAError('Test can not processed on OTHER board type devices')
-        autotest_client = autotest.Autotest(host)
-        host.reboot()
-        autotest_client.run_test(client_autotest,
-                                 exit_without_logout=True)
         if not host.run(self.CMD_CORRUPT,
                         ignore_status=True).exit_status == 0:
-             raise error.TestFail('Unable to corrupt stateful partition')
+            raise error.TestFail('Unable to corrupt stateful partition')
         host.run('sync', ignore_status=True)
         time.sleep(self._WAIT_DELAY)
         host.reboot()
@@ -49,7 +53,7 @@
             raise error.TestFail('Did not get OOBE screen after '
                                  'rebooting the device with '
                                  'corrupted statefull partition')
-        autotest_client.run_test(client_autotest,
+        autotest_client.run_test(client_test,
                                  exit_without_logout=True)
         time.sleep(self._WAIT_DELAY)
         for new_file in self.FILES_LIST:
@@ -65,3 +69,20 @@
             raise error.TestFail('An encryption key is missing after '
                                  'rebooting the device with corrupted stateful '
                                  'partition')
+
+    def run_once(self, host, client_test):
+        """This test verify that user should get OOBE after booting
+        the device with corrupted stateful partition.
+        Test fails if not able to recover the device with corrupted
+        stateful partition.
+        TPM is reset at the end to restore the system to a reasonable state.
+        """
+        if host.get_board_type() == 'OTHER':
+            raise error.TestNAError('Test can not processed on OTHER board '
+                                    'type devices')
+        autotest_client = autotest.Autotest(host)
+        host.reboot()
+        autotest_client.run_test(client_test,
+                                 exit_without_logout=True)
+        with ensure_tpm_reset(host):
+            self._test_stateful_corruption(autotest_client, host, client_test)