diff --git a/locust/core.py b/locust/core.py index e5236d328e..aa4be45d69 100644 --- a/locust/core.py +++ b/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): """ diff --git a/locust/test/mock_logging.py b/locust/test/mock_logging.py new file mode 100644 index 0000000000..13687b790c --- /dev/null +++ b/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, []) diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py index 292dc86c0f..6da4012ec9 100644 --- a/locust/test/test_runners.py +++ b/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): diff --git a/locust/test/test_zmqrpc.py b/locust/test/test_zmqrpc.py index 01b02b1ada..d6bbdde6f3 100644 --- a/locust/test/test_zmqrpc.py +++ b/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) diff --git a/locust/test/testcases.py b/locust/test/testcases.py index 1b213047ee..bf9bbce4df 100644 --- a/locust/test/testcases.py +++ b/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):