diff --git a/server/site_tests/autoupdate_DataPreserved/autoupdate_DataPreserved.py b/server/site_tests/autoupdate_DataPreserved/autoupdate_DataPreserved.py
index e1f22dd..99f5d5e 100644
--- a/server/site_tests/autoupdate_DataPreserved/autoupdate_DataPreserved.py
+++ b/server/site_tests/autoupdate_DataPreserved/autoupdate_DataPreserved.py
@@ -19,6 +19,7 @@
     def cleanup(self):
         self._save_extra_update_engine_logs(number_of_logs=2)
         super(autoupdate_DataPreserved, self).cleanup()
+        self._restore_stateful()
 
 
     def run_once(self, full_payload=True, job_repo_url=None):
diff --git a/server/site_tests/autoupdate_ForcedOOBEUpdate/autoupdate_ForcedOOBEUpdate.py b/server/site_tests/autoupdate_ForcedOOBEUpdate/autoupdate_ForcedOOBEUpdate.py
index 7201c7e..f674e64 100644
--- a/server/site_tests/autoupdate_ForcedOOBEUpdate/autoupdate_ForcedOOBEUpdate.py
+++ b/server/site_tests/autoupdate_ForcedOOBEUpdate/autoupdate_ForcedOOBEUpdate.py
@@ -32,7 +32,8 @@
             logging.debug('Canceling the in-progress update.')
             self._restart_update_engine()
         super(autoupdate_ForcedOOBEUpdate, self).cleanup()
-
+        if self._m2n:
+            self._restore_stateful()
 
     def _wait_for_reboot_after_update(self, timeout_minutes=15):
         """
@@ -136,7 +137,8 @@
             raise error.TestFail('Unknown interrupt type: %s' % interrupt)
         tpm_utils.ClearTPMOwnerRequest(self._host)
 
-        if m2n:
+        self._m2n = m2n
+        if self._m2n:
             # Provision latest stable build for the current build.
             build_name = self._get_latest_serving_stable_build()