faft: restore kernel attributes if they are changed during tests

For tests that change kernel attributes such as priority and tries
flags, we should ensure the attributes get restored when the tests end.

BUG=chromium:215787
TEST=Pass all changed tests
TEST=Prioritize kernel B and run a test. Check the active kernel is
still B after the test ends.

Change-Id: I3212f1396f1abc7473b42308145d04e051b27422
Signed-off-by: Vic (Chun-Ju) Yang <victoryang@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/185325
Reviewed-by: Yusuf Mohsinally <mohsinally@chromium.org>
diff --git a/server/site_tests/firmware_CgptState/firmware_CgptState.py b/server/site_tests/firmware_CgptState/firmware_CgptState.py
index efcd9a0..61c9ed8 100644
--- a/server/site_tests/firmware_CgptState/firmware_CgptState.py
+++ b/server/site_tests/firmware_CgptState/firmware_CgptState.py
@@ -27,6 +27,7 @@
 
 
     def run_test_step(self):
+        """Run the actual test steps."""
         # Show the client log messages in another thread.
         show_client_log = subprocess.Popen([
                 'ssh -o StrictHostKeyChecking=no '
@@ -44,13 +45,14 @@
     def initialize(self, host, cmdline_args):
         super(firmware_CgptState, self).initialize(host, cmdline_args)
         self.host = host
+        self.backup_cgpt_attributes()
         self.setup_dev_mode(dev_mode=False)
         self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
     def cleanup(self):
-        self.ensure_kernel_boot('a')
+        self.restore_cgpt_attributes()
         super(firmware_CgptState, self).cleanup()
 
 
diff --git a/server/site_tests/firmware_CgptStress/firmware_CgptStress.py b/server/site_tests/firmware_CgptStress/firmware_CgptStress.py
index 42d11da..3984140 100644
--- a/server/site_tests/firmware_CgptStress/firmware_CgptStress.py
+++ b/server/site_tests/firmware_CgptStress/firmware_CgptStress.py
@@ -21,13 +21,14 @@
         dict_args = utils.args_to_dict(cmdline_args)
         self.faft_iterations = int(dict_args.get('faft_iterations', 1))
         super(firmware_CgptStress, self).initialize(host, cmdline_args)
+        self.backup_cgpt_attributes()
         self.setup_dev_mode(dev_mode)
         self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
     def cleanup(self):
-        self.ensure_kernel_boot('a')
+        self.restore_cgpt_attributes()
         super(firmware_CgptStress, self).cleanup()
 
 
diff --git a/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py b/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py
index 66a307f..0b819de 100644
--- a/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py
+++ b/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py
@@ -41,6 +41,7 @@
     def initialize(self, host, cmdline_args, dev_mode=False):
         super(firmware_CorruptBothKernelAB, self).initialize(host, cmdline_args)
         self.backup_kernel()
+        self.backup_cgpt_attributes()
         self.setup_dev_mode(dev_mode)
         self.setup_usbkey(usbkey=True, host=False)
         self.setup_kernel('a')
@@ -48,6 +49,7 @@
 
     def cleanup(self):
         self.ensure_kernel_on_non_recovery('a')
+        self.restore_cgpt_attributes()
         self.restore_kernel()
         super(firmware_CorruptBothKernelAB, self).cleanup()
 
diff --git a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
index 0c5615d..5c1c8d2 100644
--- a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
+++ b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
@@ -18,13 +18,14 @@
     def initialize(self, host, cmdline_args, dev_mode=False):
         super(firmware_CorruptKernelA, self).initialize(host, cmdline_args)
         self.backup_kernel()
+        self.backup_cgpt_attributes()
         self.setup_dev_mode(dev_mode)
         self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
     def cleanup(self):
-        self.ensure_kernel_boot('a')
+        self.restore_cgpt_attributes()
         self.restore_kernel()
         super(firmware_CorruptKernelA, self).cleanup()
 
diff --git a/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py b/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
index 67a8e5c..6b91ff0 100644
--- a/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
+++ b/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
@@ -19,13 +19,14 @@
     def initialize(self, host, cmdline_args, dev_mode=False):
         super(firmware_CorruptKernelB, self).initialize(host, cmdline_args)
         self.backup_kernel()
+        self.backup_cgpt_attributes()
         self.setup_dev_mode(dev_mode)
         self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
     def cleanup(self):
-        self.ensure_kernel_boot('a')
+        self.restore_cgpt_attributes()
         self.restore_kernel()
         super(firmware_CorruptKernelB, self).cleanup()
 
diff --git a/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py b/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py
index 4710194..6558f2a 100644
--- a/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py
+++ b/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py
@@ -41,6 +41,7 @@
     def initialize(self, host, cmdline_args, dev_mode=False):
         super(firmware_RollbackKernel, self).initialize(host, cmdline_args)
         self.backup_kernel()
+        self.backup_cgpt_attributes()
         self.setup_dev_mode(dev_mode)
         self.setup_usbkey(usbkey=True, host=False)
         self.setup_kernel('a')
@@ -48,6 +49,7 @@
 
     def cleanup(self):
         self.ensure_kernel_on_non_recovery('a')
+        self.restore_cgpt_attributes()
         self.restore_kernel()
         super(firmware_RollbackKernel, self).cleanup()