Model/model_data/tests/test_utils.py

import logging
from io import StringIO
from unittest.mock import patch
from model_data.utils import setup_logger


class TestLogger:
    def test_setup_logger_default(self):
        log_stream = StringIO()
        handler = logging.StreamHandler(log_stream)
        logger = setup_logger()
        logger.addHandler(handler)

        # log something
        logger.info("Hello World!")

        log_stream.seek(0)
        # assert that log was written
        assert log_stream.read() == "Hello World!\n"
        # remove the handler after use
        logger.removeHandler(handler)

    @patch('logging.FileHandler')
    def test_setup_logger_file(self, mock_file_handler):
        # setup the logger
        logger = setup_logger(log_file='test.log', overwrite_handler=True)

        # assert FileHandler was called correctly
        mock_file_handler.assert_called_once_with('test.log')

        # clean up after use
        for handler in logger.handlers[:]:
            handler.close()
            logger.removeHandler(handler)

    def test_setup_logger_loglevel(self):
        log_stream = StringIO()
        handler = logging.StreamHandler(log_stream)
        logger = setup_logger(level=logging.DEBUG)
        logger.addHandler(handler)

        # log something
        logger.debug("Hello World!")

        log_stream.seek(0)
        # assert that log was written
        assert log_stream.read() == "Hello World!\n"
        # remove the handler after use
        logger.removeHandler(handler)