Using addCleanup() for restoring original configuration contents in tests.

tests/__init__.py

@@ -135,6 +135,11 @@ class TestBase(unittest.TestCase):
     def tearDown(self):
         del self.data
 
+    def restore_config(self, content):
+        """Helper for restoring configuration contents on cleanup."""
+        with open(self.britney_conf, 'w') as fp:
+            fp.write(content)
+
     def run_britney(self, args=[]):
         '''Run britney.
 

tests/test_autopkgtest.py

@@ -32,14 +32,14 @@ class TestAutoPkgTest(TestBase):
         super(TestAutoPkgTest, self).setUp()
 
         # Mofify configuration according to the test context.
-        self.old_config = None
         with open(self.britney_conf, 'r') as fp:
-            self.old_config = fp.read()
+            original_config = fp.read()
         # Disable boottests.
-        config = self.old_config.replace(
+        new_config = original_config.replace(
             'BOOTTEST_ENABLE   = yes', 'BOOTTEST_ENABLE   = no')
         with open(self.britney_conf, 'w') as fp:
-            fp.write(config)
+            fp.write(new_config)
+        self.addCleanup(self.restore_config, original_config)
 
         # fake adt-britney script
         self.adt_britney = os.path.join(
@@ -63,12 +63,6 @@ echo "$@" >> /%s/adt-britney.log ''' % self.data.path)
                                       'Conflicts': 'green'})
         self.data.add('justdata', False, {'Architecture': 'all'})
 
-    def tearDown(self):
-        """Replace the old_config."""
-        with open(self.britney_conf, 'w') as fp:
-            fp.write(self.old_config)
-        super(TestAutoPkgTest, self).tearDown()
-
     def __merge_records(self, results, history=""):
         '''Merges a list of results with records in history.
 

tests/test_boottest.py

@@ -117,17 +117,19 @@ class TestBoottestEnd2End(TestBase):
 
     def setUp(self):
         super(TestBoottestEnd2End, self).setUp()
-        self.old_config = None
+
+        # Modify shared configuration file.
         with open(self.britney_conf, 'r') as fp:
-            self.old_config = fp.read()
+            original_config = fp.read()
         # Disable autopkgtests.
-        config = self.old_config.replace(
+        new_config = original_config.replace(
             'ADT_ENABLE        = yes', 'ADT_ENABLE        = no')
         # Disable TouchManifest auto-fetching.
-        config = config.replace(
+        new_config = new_config.replace(
             'BOOTTEST_FETCH    = yes', 'BOOTTEST_FETCH    = no')
         with open(self.britney_conf, 'w') as fp:
-            fp.write(config)
+            fp.write(new_config)
+        self.addCleanup(self.restore_config, original_config)
 
         self.data.add('libc6', False, {'Architecture': 'armhf'}),
 
@@ -146,12 +148,6 @@ class TestBoottestEnd2End(TestBase):
             'pyqt5:armhf 1.0',
         ])
 
-    def tearDown(self):
-        """ Replace the old_config. """
-        with open(self.britney_conf, 'w') as fp:
-            fp.write(self.old_config)
-        super(TestBoottestEnd2End, self).tearDown()
-
     def create_manifest(self, lines):
         """Create a manifest for this britney run context."""
         path = os.path.join(