Use a mocked logging handler while running test logging messages

locust/core.py

@@ -23,8 +23,6 @@
 from .runners import STATE_CLEANUP, LOCUST_STATE_RUNNING, LOCUST_STATE_STOPPING, LOCUST_STATE_WAITING
 from .util import deprecation
 
-logger = logging.getLogger(__name__)
-
 
 def task(weight=1):
     """

locust/test/mock_logging.py

@@ -0,0 +1,17 @@
+import logging
+
+
+class MockedLoggingHandler(logging.Handler):
+    debug = []
+    warning = []
+    info = []
+    error = []
+
+    def emit(self, record):
+        getattr(self.__class__, record.levelname.lower()).append(record.getMessage())
+
+    @classmethod
+    def reset(cls):
+        for attr in dir(cls):
+            if isinstance(getattr(cls, attr), list):
+                setattr(cls, attr, [])

locust/test/test_runners.py

@@ -136,13 +136,14 @@ def trigger(self):
 
 class TestMasterRunner(LocustTestCase):
     def setUp(self):
+        super(TestMasterRunner, self).setUp()
         global_stats.reset_all()
         self._slave_report_event_handlers = [h for h in events.slave_report._handlers]
         self.options = mocked_options()
-
 
     def tearDown(self):
         events.slave_report._handlers = self._slave_report_event_handlers
+        super(TestMasterRunner, self).tearDown()
 
     def test_slave_connect(self):
         class MyTestLocust(Locust):
@@ -488,7 +489,7 @@ def test_message_serialize(self):
         self.assertEqual(msg.data, rebuilt.data)
         self.assertEqual(msg.node_id, rebuilt.node_id)
 
-class TestStopTimeout(unittest.TestCase):
+class TestStopTimeout(LocustTestCase):
     def test_stop_timeout(self):
         short_time = 0.05
         class MyTaskSet(TaskSet):

locust/test/test_zmqrpc.py

@@ -2,16 +2,19 @@
 from time import sleep
 import zmq
 from locust.rpc import zmqrpc, Message
+from locust.test.testcases import LocustTestCase
 
 
-class ZMQRPC_tests(unittest.TestCase):
+class ZMQRPC_tests(LocustTestCase):
     def setUp(self):
+        super(ZMQRPC_tests, self).setUp()
         self.server = zmqrpc.Server('*', 0)
         self.client = zmqrpc.Client('localhost', self.server.port, 'identity')
 
     def tearDown(self):
         self.server.socket.close()
         self.client.socket.close()
+        super(ZMQRPC_tests, self).tearDown()
 
     def test_constructor(self):
         self.assertEqual(self.server.socket.getsockopt(zmq.TCP_KEEPALIVE), 1)

locust/test/testcases.py

@@ -1,4 +1,5 @@
 import base64
+import logging
 import random
 import sys
 import unittest
@@ -13,7 +14,9 @@
                    send_file, stream_with_context)
 
 from locust import events
+from locust.log import console_logger
 from locust.stats import global_stats
+from locust.test.mock_logging import MockedLoggingHandler
 
 
 app = Flask(__name__)
@@ -138,10 +141,29 @@ def setUp(self):
             # ResourceWarning doesn't exist in Python 2, but since the warning only appears
             # on Python 3 we don't need to mock it. Instead we can happily ignore the exception
             pass
+
+        # set up mocked logging handler
+        self._logger_class = MockedLoggingHandler()
+        self._logger_class.setLevel(logging.INFO)
+        console_logger.propagate = True
+        self._root_log_handlers = [h for h in logging.root.handlers]
+        self._console_log_handlers = [h for h in console_logger.handlers]
+        [logging.root.removeHandler(h) for h in logging.root.handlers]
+        [console_logger.removeHandler(h) for h in console_logger.handlers]
+        logging.root.addHandler(self._logger_class)
+        logging.root.setLevel(logging.INFO)
+        self.mocked_log = MockedLoggingHandler
 
     def tearDown(self):
         for event, handlers in six.iteritems(self._event_handlers):
             event._handlers = handlers
+
+        # restore logging class
+        logging.root.removeHandler(self._logger_class)
+        [logging.root.addHandler(h) for h in self._root_log_handlers]
+        [console_logger.addHandler(h) for h in self._console_log_handlers]
+        self.mocked_log.reset()
+        console_logger.propagate = False
 
 
 class WebserverTestCase(LocustTestCase):