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