tests: Ensure additional log handlers are removed

In some places handlers added using logger.addHandler were not removed.
Some tests manually added MockLoggingHandler instead of using `with` .

This commit fixes those problems.

Author: Lorak-mmk

tests/integration/__init__.py

@@ -896,7 +896,7 @@ def __enter__(self):
         return self
 
     def __exit__(self, *args):
-        pass
+        self.logger.removeHandler(self)
 
 
 class BasicExistingKeyspaceUnitTestCase(BasicKeyspaceUnitTestCase):

tests/integration/cqlengine/management/test_management.py

@@ -360,18 +360,15 @@ def test_sync_warnings(self):
 
         @test_category object_mapper
         """
-        mock_handler = MockLoggingHandler()
-        logger = logging.getLogger(management.__name__)
-        logger.addHandler(mock_handler)
-        sync_table(BaseInconsistent)
-        sync_table(ChangedInconsistent)
-        self.assertTrue('differing from the model type' in mock_handler.messages.get('warning')[0])
-        if CASSANDRA_VERSION >= Version('2.1'):
-            sync_type(DEFAULT_KEYSPACE, BaseInconsistentType)
-            mock_handler.reset()
-            sync_type(DEFAULT_KEYSPACE, ChangedInconsistentType)
-            self.assertTrue('differing from the model user type' in mock_handler.messages.get('warning')[0])
-        logger.removeHandler(mock_handler)
+        with MockLoggingHandler().set_module_name(management.__name__) as mock_handler:
+            sync_table(BaseInconsistent)
+            sync_table(ChangedInconsistent)
+            self.assertTrue('differing from the model type' in mock_handler.messages.get('warning')[0])
+            if CASSANDRA_VERSION >= Version('2.1'):
+                sync_type(DEFAULT_KEYSPACE, BaseInconsistentType)
+                mock_handler.reset()
+                sync_type(DEFAULT_KEYSPACE, ChangedInconsistentType)
+                self.assertTrue('differing from the model user type' in mock_handler.messages.get('warning')[0])
 
 
 class TestIndexSetModel(Model):

tests/integration/simulacron/test_cluster.py

@@ -88,23 +88,20 @@ class DuplicateRpcTest(SimulacronCluster):
     connect = False
 
     def test_duplicate(self):
-        mock_handler = MockLoggingHandler()
-        logger = logging.getLogger(cassandra.cluster.__name__)
-        logger.addHandler(mock_handler)
-        address_column = "native_transport_address" if DSE_VERSION and DSE_VERSION > Version("6.0") else "rpc_address"
-        rows = [
-            {"peer": "127.0.0.1", "data_center": "dc", "host_id": "dontcare1", "rack": "rack1",
-             "release_version": "3.11.4", address_column: "127.0.0.1", "schema_version": "dontcare", "tokens": "1"},
-            {"peer": "127.0.0.2", "data_center": "dc", "host_id": "dontcare2", "rack": "rack1",
-             "release_version": "3.11.4", address_column: "127.0.0.2", "schema_version": "dontcare", "tokens": "2"},
-        ]
-        prime_query(ControlConnection._SELECT_PEERS, rows=rows)
-
-        cluster = Cluster(protocol_version=PROTOCOL_VERSION, compression=False)
-        session = cluster.connect(wait_for_all_pools=True)
-
-        warnings = mock_handler.messages.get("warning")
-        self.assertEqual(len(warnings), 1)
-        self.assertTrue('multiple hosts with the same endpoint' in warnings[0])
-        logger.removeHandler(mock_handler)
-        cluster.shutdown()
+        with MockLoggingHandler().set_module_name(cassandra.cluster.__name__) as mock_handler:
+            address_column = "native_transport_address" if DSE_VERSION and DSE_VERSION > Version("6.0") else "rpc_address"
+            rows = [
+                {"peer": "127.0.0.1", "data_center": "dc", "host_id": "dontcare1", "rack": "rack1",
+                "release_version": "3.11.4", address_column: "127.0.0.1", "schema_version": "dontcare", "tokens": "1"},
+                {"peer": "127.0.0.2", "data_center": "dc", "host_id": "dontcare2", "rack": "rack1",
+                "release_version": "3.11.4", address_column: "127.0.0.2", "schema_version": "dontcare", "tokens": "2"},
+            ]
+            prime_query(ControlConnection._SELECT_PEERS, rows=rows)
+
+            cluster = Cluster(protocol_version=PROTOCOL_VERSION, compression=False)
+            session = cluster.connect(wait_for_all_pools=True)
+
+            warnings = mock_handler.messages.get("warning")
+            self.assertEqual(len(warnings), 1)
+            self.assertTrue('multiple hosts with the same endpoint' in warnings[0])
+            cluster.shutdown()

tests/integration/standard/test_query.py

@@ -508,6 +508,9 @@ def setUp(self):
         self.mock_handler = MockLoggingHandler()
         logger = logging.getLogger(cluster.__name__)
         logger.addHandler(self.mock_handler)
+    
+    def tearDown(self):
+        logger.removeHandler(self.mock_handler)
 
     def test_prepare_on_all_hosts(self):
         """
@@ -1562,28 +1565,27 @@ def test_reprepare_after_host_is_down(self):
 
         @test_category query
         """
-        mock_handler = MockLoggingHandler()
-        logger = logging.getLogger(cluster.__name__)
-        logger.addHandler(mock_handler)
-        get_node(1).stop(wait=True, gently=True, wait_other_notice=True)
+        with MockLoggingHandler().set_module_name(cluster.__name__) as mock_handler:
+            get_node(1).stop(wait=True, gently=True, wait_other_notice=True)
 
-        only_first = ExecutionProfile(load_balancing_policy=WhiteListRoundRobinPolicy(["127.0.0.1"]))
-        self.cluster.add_execution_profile("only_first", only_first)
+            only_first = ExecutionProfile(load_balancing_policy=WhiteListRoundRobinPolicy(["127.0.0.1"]))
+            self.cluster.add_execution_profile("only_first", only_first)
 
-        query = "SELECT v from {} WHERE k = ?".format(self.table_name)
-        prepared_statement = self.session.prepare(query, keyspace=self.ks_name)
-        prepared_statement_alternative = self.session.prepare(query, keyspace=self.alternative_ks)
+            query = "SELECT v from {} WHERE k = ?".format(self.table_name)
+            prepared_statement = self.session.prepare(query, keyspace=self.ks_name)
+            prepared_statement_alternative = self.session.prepare(query, keyspace=self.alternative_ks)
 
-        get_node(1).start(wait_for_binary_proto=True, wait_other_notice=True)
+            get_node(1).start(wait_for_binary_proto=True, wait_other_notice=True)
 
-        # We wait for cluster._prepare_all_queries to be called
-        time.sleep(5)
-        self.assertEqual(1, mock_handler.get_message_count('debug', 'Preparing all known prepared statements'))
-        results = self.session.execute(prepared_statement, (1,), execution_profile="only_first")
-        self.assertEqual(results[0], (1, ))
+            # We wait for cluster._prepare_all_queries to be called
+            time.sleep(5)
+            self.assertEqual(1, mock_handler.get_message_count('debug', 'Preparing all known prepared statements'))
+            
+            results = self.session.execute(prepared_statement, (1,), execution_profile="only_first")
+            self.assertEqual(results[0], (1, ))
 
-        results = self.session.execute(prepared_statement_alternative, (2,), execution_profile="only_first")
-        self.assertEqual(results[0], (2, ))
+            results = self.session.execute(prepared_statement_alternative, (2,), execution_profile="only_first")
+            self.assertEqual(results[0], (2, ))
 
     def test_prepared_not_found(self):
         """

tests/integration/upgrade/__init__.py

@@ -78,6 +78,10 @@ def setUpClass(cls):
         cls.logger_handler = MockLoggingHandler()
         logger = logging.getLogger(cluster.__name__)
         logger.addHandler(cls.logger_handler)
+    
+    @classmethod
+    def tearDownClass(cls):
+        logger.removeHandler(cls.logger_handler)
 
     def _upgrade_step_setup(self):
         """