From bfc6fc19aa30fdf9ec4fa29c77fb81f672c95035 Mon Sep 17 00:00:00 2001
From: Jim Porter <itsjimporter@gmail.com>
Date: Thu, 7 Dec 2023 12:51:15 -0800
Subject: [PATCH] Use `TestCase.addCleanup` where possible

---
 test/unit/builtins/test_link.py | 17 ++++++++---------
 test/unit/test_driver.py        |  7 +++----
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/test/unit/builtins/test_link.py b/test/unit/builtins/test_link.py
index 1bce533d..02d4bcf8 100644
--- a/test/unit/builtins/test_link.py
+++ b/test/unit/builtins/test_link.py
@@ -1,16 +1,19 @@
 from unittest import mock
 
+from .. import make_env
 from .common import AttrDict, BuiltinTest
+
+from bfg9000 import file_types, options as opts
 from bfg9000.backends.make import syntax as make
 from bfg9000.backends.msbuild.solution import Solution
 from bfg9000.backends.ninja import syntax as ninja
 from bfg9000.builtins import (compile, default, link, packages,  # noqa: F401
                               project)
-from bfg9000 import file_types, options as opts
 from bfg9000.environment import LibraryMode
 from bfg9000.iterutils import listify, unlistify
 from bfg9000.packages import CommonPackage
 from bfg9000.path import Path, Root
+from bfg9000.tools.msvc import MsvcBuilder
 
 
 class LinkTest(BuiltinTest):
@@ -938,18 +941,14 @@ def test_extra_deps(self):
 
 class TestMsbuildBackend(BuiltinTest):
     def setUp(self):
-        from .. import make_env
         self.env = make_env('winnt', clear_variables=True,
                             variables={'CXX': 'nonexist'})
         self.build, self.context = self._make_context(self.env)
 
-        from bfg9000.tools.msvc import MsvcBuilder
-        self.patch_builder = mock.patch('bfg9000.tools.c_family._builders',
-                                        (MsvcBuilder,))
-        self.patch_builder.start()
-
-    def tearDown(self):
-        self.patch_builder.stop()
+        patch_builder = mock.patch('bfg9000.tools.c_family._builders',
+                                   (MsvcBuilder,))
+        patch_builder.start()
+        self.addCleanup(patch_builder.stop)
 
     def assertSubdict(self, actual, expected):
         subdict = {k: v for k, v in actual.items() if k in expected}
diff --git a/test/unit/test_driver.py b/test/unit/test_driver.py
index 57cd9421..4b9ac82a 100644
--- a/test/unit/test_driver.py
+++ b/test/unit/test_driver.py
@@ -75,11 +75,10 @@ def setUp(self):
         self.logger = log.getLogger('bfg9000.test.unit')
         self.logger.propagate = False
         log._init_logging(self.logger, debug=False, stream=self.stream)
-        self.patch_logger = mock.patch('bfg9000.driver.logger', self.logger)
-        self.patch_logger.start()
 
-    def tearDown(self):
-        self.patch_logger.stop()
+        patch_logger = mock.patch('bfg9000.driver.logger', self.logger)
+        patch_logger.start()
+        self.addCleanup(patch_logger.stop)
 
     def test_message(self):
         try: