local_agent/tests/unit_tests/test_logger.py

# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for logger."""
import unittest
from unittest import mock
import logging

from local_agent import logger

_FAKE_LOG_FILE = 'fake-log-file'


class LoggerTest(unittest.TestCase):
    """Unit test for logger module."""

    @mock.patch.object(logging.handlers, 'RotatingFileHandler')
    def test_01_create_handler_log_file_not_none(
        self, mock_rotating_file_handler):
        """Verifies create_handler creates logging handler for log file."""
        mock_handler = mock.Mock()
        mock_rotating_file_handler.return_value = mock_handler

        returned_handler = logger.create_handler(log_file=_FAKE_LOG_FILE)

        self.assertEqual(returned_handler, mock_handler)
        mock_handler.setLevel.assert_called_once()
        mock_handler.setFormatter.assert_called_once()

    @mock.patch.object(logging, 'StreamHandler')
    def test_01_create_handler_no_log_file(self, mock_stream_handler):
        """Verifies create_handler creates logging stream handler."""
        mock_handler = mock.Mock()
        mock_stream_handler.return_value = mock_handler

        returned_handler = logger.create_handler()

        self.assertEqual(returned_handler, mock_handler)
        mock_handler.setLevel.assert_called_once()
        mock_handler.setFormatter.assert_called_once()

    @mock.patch.object(logger, 'logger')
    @mock.patch.object(logger, 'create_handler')
    def test_02_add_file_handler(self, mock_create_handler, mock_logger):
        """Verifies add_file_handler adds handler for log file."""
        mock_handler = mock.Mock()
        mock_create_handler.return_value = mock_handler

        logger.add_file_handler(log_file=_FAKE_LOG_FILE)

        mock_create_handler.assert_called_once()
        self.assertEqual(logger.log_handler_map[_FAKE_LOG_FILE], mock_handler)
        mock_logger.addHandler.assert_called_once_with(mock_handler)

    @mock.patch.object(logger, 'create_handler')
    def test_02_add_file_handler_already_exists(
        self, mock_create_handler):
        """Verifies add_file_handler does nothing for already existing file."""
        logger.log_handler_map[_FAKE_LOG_FILE] = mock.Mock()

        logger.add_file_handler(log_file=_FAKE_LOG_FILE)

        self.assertEqual(0, mock_create_handler.call_count)

    @mock.patch.object(logger, 'logger')
    def test_03_remove_file_handler(self, mock_logger):
        """Verifies remove_file_handler removes handler for log file."""
        mock_handler = mock.Mock()
        logger.log_handler_map[_FAKE_LOG_FILE] = mock_handler

        logger.remove_file_handler(log_file=_FAKE_LOG_FILE)

        self.assertNotIn(_FAKE_LOG_FILE, logger.log_handler_map)
        mock_logger.removeHandler.assert_called_once_with(mock_handler)

    @mock.patch.object(logger, 'logger')
    def test_03_remove_file_handler_not_exists(self, mock_logger):
        """Verifies remove_file_handler does nothing for non-existed handler."""
        logger.log_handler_map.clear()

        logger.remove_file_handler(log_file=_FAKE_LOG_FILE)

        self.assertEqual(0, mock_logger.removeHandler.call_count)

    @mock.patch.object(logger, 'create_handler')
    @mock.patch.object(logger, 'logger')
    def test_04_setup_logger(self, mock_logger, mock_create_handler):
        """Verifies setup_logger on success."""
        logger.setup_logger()

        mock_logger.setLevel.assert_called_once()
        self.assertEqual(2, mock_create_handler.call_count)
        self.assertEqual(2, mock_logger.addHandler.call_count)

    def test_05_get_logger(self):
        """Verifies get_logger gets local agent logger."""
        local_agent_logger = logger.get_logger()
        self.assertEqual(logger.logger, local_agent_logger)


if __name__ == '__main__':
    unittest.main(failfast=True)