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
5import logging
6import logging.handlers
7import mox
8import multiprocessing
9import tempfile
10import time
11import os
12import unittest
13
14import log_socket_server
15
16
17class TestLogSocketServer(mox.MoxTestBase):
18    """Test LogSocketServer can start and save logs to a local file.
19    """
20
21
22    def log_call(self, value, port):
23        """Method to be called in a new process to log to a socket server.
24
25        @param value: Data to be logged.
26        @param port: Port used by log socket server.
27        """
28        logging.getLogger().handlers = []
29        socketHandler = logging.handlers.SocketHandler('localhost', port)
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
44        port = log_socket_server.LogSocketServer.port
45        for i in range(process_number):
46            process = multiprocessing.Process(target=self.log_call,
47                                              args=(i, port))
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))
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))
68        os.remove(log_filename)
69
70
71if __name__ == "__main__":
72    unittest.main()
73