Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 1 | # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. |
| 2 | # Use of this source code is governed by a BSD-style license that can be |
| 3 | # found in the LICENSE file. |
| 4 | |
| 5 | import logging |
| 6 | import logging.handlers |
| 7 | import mox |
| 8 | import multiprocessing |
| 9 | import tempfile |
| 10 | import time |
| 11 | import os |
| 12 | import unittest |
| 13 | |
| 14 | import log_socket_server |
| 15 | |
| 16 | |
| 17 | class TestLogSocketServer(mox.MoxTestBase): |
| 18 | """Test LogSocketServer can start and save logs to a local file. |
| 19 | """ |
| 20 | |
| 21 | |
Dan Shi | cd0a01d | 2014-06-11 20:53:35 -0700 | [diff] [blame] | 22 | def log_call(self, value, port): |
Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 23 | """Method to be called in a new process to log to a socket server. |
| 24 | |
| 25 | @param value: Data to be logged. |
Dan Shi | cd0a01d | 2014-06-11 20:53:35 -0700 | [diff] [blame] | 26 | @param port: Port used by log socket server. |
Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 27 | """ |
| 28 | logging.getLogger().handlers = [] |
Dan Shi | cd0a01d | 2014-06-11 20:53:35 -0700 | [diff] [blame] | 29 | socketHandler = logging.handlers.SocketHandler('localhost', port) |
Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 30 | logging.getLogger().addHandler(socketHandler) |
| 31 | logging.getLogger().level = logging.INFO |
| 32 | logging.info(value) |
| 33 | |
| 34 | |
| 35 | def testMultiProcessLoggingSuccess(self): |
| 36 | """Test log can be saved from multiple processes.""" |
| 37 | # Start log TCP server. |
| 38 | super(TestLogSocketServer, self).setUp() |
| 39 | log_filename = tempfile.mktemp(suffix='_log_server') |
| 40 | log_socket_server.LogSocketServer.start(filename=log_filename, |
| 41 | level=logging.INFO) |
| 42 | processes = [] |
| 43 | process_number = 10 |
Dan Shi | cd0a01d | 2014-06-11 20:53:35 -0700 | [diff] [blame] | 44 | port = log_socket_server.LogSocketServer.port |
Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 45 | for i in range(process_number): |
Dan Shi | cd0a01d | 2014-06-11 20:53:35 -0700 | [diff] [blame] | 46 | process = multiprocessing.Process(target=self.log_call, |
| 47 | args=(i, port)) |
Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 48 | process.start() |
| 49 | processes.append(process) |
| 50 | |
| 51 | for process in processes: |
| 52 | process.join() |
| 53 | |
| 54 | # Wait for TCP server to finish processing data. If process_number is |
| 55 | # increased, the wait time should be increased to avoid test flaky. |
| 56 | time.sleep(1) |
| 57 | log_socket_server.LogSocketServer.stop() |
| 58 | |
| 59 | # Read log to confirm all logs are written to file. |
| 60 | num_lines = sum(1 for line in open(log_filename)) |
Dan Shi | 0ceff75 | 2014-09-04 12:23:07 -0700 | [diff] [blame] | 61 | if process_number != num_lines: |
| 62 | logging.warn('Not all log messages were written to file %s. ' |
| 63 | 'Expected number of logs: %s, Logs found in file: %s', |
| 64 | log_filename, process_number, num_lines) |
| 65 | self.assertNotEqual(0, num_lines, 'No log message was written to file ' |
| 66 | '%s. Number of logs tried: %s.' % |
| 67 | (log_filename, process_number)) |
Dan Shi | 1a34c36 | 2014-04-11 16:37:04 -0700 | [diff] [blame] | 68 | os.remove(log_filename) |
| 69 | |
| 70 | |
| 71 | if __name__ == "__main__": |
Dan Shi | cd0a01d | 2014-06-11 20:53:35 -0700 | [diff] [blame] | 72 | unittest.main() |