ComfyAssets-ComfyUI_PromptManager/tests/test_utils_logging_config.py

"""PromptManagerLogger: memory buffer, file rotation, level changes, truncation."""

import logging
import os
import pathlib
import sys
import tempfile
import threading
import unittest
import unittest.mock as mock

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils import logging_config
from utils.logging_config import (
    MemoryBufferHandler,
    PromptManagerLogger,
    get_logger,
    get_logger_manager,
)


class IsolatedLoggerTestCase(unittest.TestCase):
    """A logger manager whose log directory is a temporary folder."""

    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        package = os.path.join(self.tmp.name, "pkg")
        os.makedirs(os.path.join(package, "utils"))
        self.log_dir = os.path.join(package, "logs")
        fake_file = os.path.join(package, "utils", "logging_config.py")
        with mock.patch.object(logging_config, "__file__", fake_file):
            self.manager = object.__new__(PromptManagerLogger)  # bypass singleton
            self.manager.__init__()
        self.manager.update_config({"console_logging": False})
        # Drop the "Updated logging configuration" entry.
        self.manager._log_buffer.clear()
        # Other tests leave daemon threads (image monitor, watchdog) logging
        # under "prompt_manager.*"; only records from this test's thread may
        # reach the buffer or the rotating file (a stray line after a truncate
        # would roll the file over and resurrect the backups). update_config
        # rebuilds the handlers, so the filter is re-applied after every call.
        self._filter_to_this_thread()
        original_update = self.manager.update_config

        def update_and_refilter(new_config):
            original_update(new_config)
            self._filter_to_this_thread()

        self.manager.update_config = update_and_refilter
        self.logger = logging.getLogger("prompt_manager.coverage_test")
        self.addCleanup(self._restore_global_logging)

    def _filter_to_this_thread(self):
        this_thread = threading.get_ident()
        for handler in self.manager.logger.handlers:
            handler.addFilter(lambda record: record.thread == this_thread)

    def _restore_global_logging(self):
        for handler in logging.getLogger("prompt_manager").handlers[:]:
            handler.close()
        get_logger_manager()._setup_loggers()
        self.tmp.cleanup()

    def _rotate(self, max_file_size=400, backup_count=2, messages=30):
        self.manager.update_config(
            {"max_file_size": max_file_size, "backup_count": backup_count}
        )
        for i in range(messages):
            self.logger.info("message %03d %s", i, "x" * 80)


class TestSingleton(unittest.TestCase):
    def test_manager_and_loggers_are_shared(self):
        self.assertIs(PromptManagerLogger(), PromptManagerLogger())
        self.assertIs(get_logger_manager(), get_logger_manager())
        self.assertIs(
            get_logger("prompt_manager.x"), logging.getLogger("prompt_manager.x")
        )
        self.assertEqual(get_logger().name, "prompt_manager")

    def test_init_is_a_noop_on_the_existing_instance(self):
        manager = get_logger_manager()
        before = manager.config
        manager.__init__()
        self.assertIs(manager.config, before)


class TestMemoryBuffer(IsolatedLoggerTestCase):
    def test_records_are_buffered_most_recent_first(self):
        self.logger.info("first")
        self.logger.warning("second")
        self.logger.error("third")
        recent = self.manager.get_recent_logs(limit=2)
        self.assertEqual([r["message"] for r in recent], ["third", "second"])
        entry = recent[0]
        self.assertEqual(entry["level"], "ERROR")
        self.assertEqual(entry["logger"], self.logger.name)
        self.assertIn("third", entry["formatted"])
        for key in ("timestamp", "module", "filename", "lineno", "thread", "process"):
            self.assertIn(key, entry)

    def test_level_filter(self):
        self.logger.info("info")
        self.logger.warning("warn")
        self.logger.error("err")
        self.assertEqual(
            [r["message"] for r in self.manager.get_recent_logs(level="warning")],
            ["err", "warn"],
        )
        self.assertEqual(len(self.manager.get_recent_logs(level="NOPE")), 3)

    def test_buffer_is_capped_at_the_configured_size(self):
        for i in range(1005):
            self.logger.info("entry %d", i)
        stats = self.manager.get_log_stats()
        self.assertEqual(stats["buffer_count"], 1000)
        self.assertEqual(stats["level_counts"], {"INFO": 1000})
        self.assertEqual(stats["current_level"], "INFO")
        self.assertEqual(
            os.path.normcase(os.path.realpath(stats["log_directory"])),
            os.path.normcase(os.path.realpath(self.log_dir)),
        )
        self.assertEqual(
            self.manager.get_recent_logs(limit=1)[0]["message"], "entry 1004"
        )

    def test_handler_swallows_buffer_errors(self):
        manager = mock.Mock()
        manager.add_to_buffer.side_effect = RuntimeError("boom")
        handler = MemoryBufferHandler(manager)
        record = logging.LogRecord("x", logging.INFO, __file__, 1, "msg", (), None)
        handler.emit(record)  # must not raise
        manager.add_to_buffer.assert_called_once()


class TestFilesAndRotation(IsolatedLoggerTestCase):
    def test_rotation_creates_backups_and_lists_them_newest_first(self):
        self._rotate()
        files = self.manager.get_log_files()
        names = sorted(f["filename"] for f in files)
        self.assertEqual(
            names,
            ["prompt_manager.log", "prompt_manager.log.1", "prompt_manager.log.2"],
        )
        main = [f for f in files if f["is_main"]]
        self.assertEqual(len(main), 1)
        self.assertTrue(all(f["size"] > 0 for f in files if not f["is_main"]))
        modified = [f["modified"] for f in files]
        self.assertEqual(modified, sorted(modified, reverse=True))
        self.assertEqual(self.manager.get_log_stats()["log_files_count"], 3)

    def test_read_log_file_tail_and_whole(self):
        self._rotate()
        tail = self.manager.read_log_file("prompt_manager.log.1", lines=1)
        self.assertEqual(len(tail), 1)
        whole = self.manager.read_log_file("prompt_manager.log.1", lines=0)
        self.assertGreater(len(whole), 1)
        self.assertTrue(whole[-1].endswith(tail[0]))

    def test_read_log_file_rejects_missing_and_outside_paths(self):
        with self.assertRaises(FileNotFoundError):
            self.manager.read_log_file("missing.log")
        outside = os.path.join(self.tmp.name, "pkg", "outside.log")
        with open(outside, "w", encoding="utf-8") as f:
            f.write("secret\n")
        with self.assertRaises(ValueError):
            self.manager.read_log_file(os.path.join("..", "outside.log"))

    def test_read_log_file_reraises_read_errors(self):
        self.logger.info("something")
        with mock.patch("builtins.open", side_effect=OSError("disk")):
            with self.assertRaises(OSError):
                self.manager.read_log_file("prompt_manager.log")

    def test_unreadable_files_are_skipped_in_the_listing(self):
        self.logger.info("something")
        original_stat = pathlib.Path.stat

        def stat_failing_for_logs(path, *args, **kwargs):
            if path.suffix == ".log":
                raise OSError("denied")
            return original_stat(path, *args, **kwargs)

        with mock.patch.object(pathlib.Path, "stat", autospec=True) as stat:
            stat.side_effect = stat_failing_for_logs
            self.assertEqual(self.manager.get_log_files(), [])

    def test_truncate_does_not_roll_the_log_over_while_deleting_backups(self):
        # With a file limit smaller than one record, any log line written
        # while the backups are being deleted rolls the main log over and
        # resurrects prompt_manager.log.1. The handler's own stream must also
        # be reset, or its stale offset reports the old size, rolls over on the
        # next record anyway and pads the file with zeros on POSIX.
        self._rotate(max_file_size=400, backup_count=2, messages=10)
        result = self.manager.truncate_logs()
        self.assertEqual(result["errors"], [])
        self.logger.info("after truncate")
        leftovers = sorted(p.name for p in pathlib.Path(self.log_dir).iterdir())
        contents = {
            name: pathlib.Path(self.log_dir, name).read_bytes() for name in leftovers
        }
        self.assertEqual(leftovers, ["prompt_manager.log"], contents)
        main_log = pathlib.Path(self.log_dir, "prompt_manager.log").read_bytes()
        self.assertNotIn(b"\x00", main_log)
        self.assertIn(b"after truncate", main_log)

    def test_truncate_clears_main_log_backups_and_buffer(self):
        self._rotate()
        result = self.manager.truncate_logs()
        self.assertEqual(result["errors"], [])
        self.assertEqual(
            sorted(result["truncated"]),
            ["prompt_manager.log", "prompt_manager.log.1", "prompt_manager.log.2"],
        )
        self.assertFalse(
            os.path.exists(os.path.join(self.log_dir, "prompt_manager.log.1"))
        )
        # Only the "Truncated log file" lines written after clearing remain
        self.assertEqual(
            [
                r["message"]
                for r in self.manager.get_recent_logs()
                if "message 0" in r["message"]
            ],
            [],
        )

    def test_truncate_reports_files_it_cannot_remove(self):
        self._rotate()
        with mock.patch.object(pathlib.Path, "unlink", side_effect=OSError("busy")):
            result = self.manager.truncate_logs()
        self.assertEqual(len(result["errors"]), 2)
        self.assertEqual(result["truncated"], ["prompt_manager.log"])


class TestConfiguration(IsolatedLoggerTestCase):
    def test_in_flight_record_cannot_reopen_retired_file_handler(self) -> None:
        handler = next(
            h
            for h in self.manager.logger.handlers
            if isinstance(h, logging.FileHandler)
        )
        stream = handler.stream
        selected = threading.Event()
        resume = threading.Event()
        errors = []
        message = "in-flight record selected before file logging disabled"

        def pause_before_handler_lock(record: logging.LogRecord) -> bool:
            if record.getMessage() == message:
                selected.set()
                if not resume.wait(5):
                    raise RuntimeError("logging thread was not released")
            return True

        def log_from_background_thread() -> None:
            try:
                self.manager.logger.info(message)
            except BaseException as error:
                errors.append(error)

        # Allow this record through the fixture's normal main-thread filter.
        handler.filters.clear()
        handler.addFilter(pause_before_handler_lock)
        worker = threading.Thread(target=log_from_background_thread, daemon=True)
        worker.start()
        try:
            self.assertTrue(selected.wait(5), "logger did not select the file handler")
            self.assertFalse(stream.closed)
            self.manager.update_config({"file_logging": False})
            self.assertNotIn(handler, self.manager.logger.handlers)
            self.assertTrue(stream.closed)
            self.assertIsNone(handler.stream)

            resume.set()
            worker.join(5)
            self.assertFalse(worker.is_alive(), "logging thread did not finish")
            self.assertEqual(errors, [])
            self.assertIsNone(handler.stream)
            self.assertTrue(stream.closed)
            self.assertFalse(self.manager.config["file_logging"])
            self.assertFalse(
                any(
                    isinstance(h, logging.FileHandler)
                    for h in self.manager.logger.handlers
                )
            )
            self.assertNotIn(
                message,
                pathlib.Path(self.log_dir, "prompt_manager.log").read_text(),
            )
        finally:
            resume.set()
            worker.join(5)
            handler.close()

    def test_rebuild_closes_replaced_file_stream(self) -> None:
        old_handler = next(
            h
            for h in self.manager.logger.handlers
            if isinstance(h, logging.FileHandler)
        )
        stream = old_handler.stream
        self.addCleanup(old_handler.close)
        self.assertFalse(stream.closed)

        self.manager.update_config({"file_logging": False})

        self.assertTrue(stream.closed)
        self.assertIsNone(old_handler.stream)
        self.assertNotIn(old_handler, self.manager.logger.handlers)
        log_file = pathlib.Path(self.log_dir, "prompt_manager.log")
        renamed = log_file.rename(log_file.with_suffix(".retired"))
        renamed.unlink()
        self.assertFalse(renamed.exists())

    def test_repeated_rebuilds_retire_previous_handlers(self) -> None:
        for change in (
            {"max_file_size": 4096},
            {"backup_count": 3},
            {"console_logging": True},
            {"console_logging": False},
        ):
            with self.subTest(change=change):
                previous = self.manager.logger.handlers[:]
                for handler in previous:
                    self.addCleanup(handler.close)
                old_file = next(
                    h for h in previous if isinstance(h, logging.FileHandler)
                )
                stream = old_file.stream
                self.assertFalse(stream.closed)
                console_streams = [
                    h.stream for h in previous if type(h) is logging.StreamHandler
                ]

                self.manager.update_config(change)

                self.assertTrue(stream.closed)
                self.assertIsNone(old_file.stream)
                self.assertTrue(all(h._closed for h in previous))
                self.assertTrue(all(not s.closed for s in console_streams))
                self.assertTrue(
                    all(h not in self.manager.logger.handlers for h in previous)
                )
                current_files = [
                    h
                    for h in self.manager.logger.handlers
                    if isinstance(h, logging.FileHandler)
                ]
                self.assertEqual(len(current_files), 1)
                self.assertFalse(current_files[0].stream.closed)
                message = f"usable replacement after {change}"
                self.logger.info(message)
                self.assertIn(
                    message,
                    pathlib.Path(self.log_dir, "prompt_manager.log").read_text(),
                )
                self.assertEqual(
                    self.manager.get_recent_logs(limit=1)[0]["message"], message
                )

    def test_rebuild_does_not_close_external_handler(self) -> None:
        external_path = pathlib.Path(self.tmp.name, "external.log")
        external_handler = logging.FileHandler(
            external_path, mode="w", encoding="utf-8"
        )
        self.addCleanup(external_handler.close)
        other_logger = logging.Logger("external_shared_sink", level=logging.INFO)
        other_logger.addHandler(external_handler)
        self.addCleanup(other_logger.removeHandler, external_handler)
        self.manager.logger.addHandler(external_handler)

        with mock.patch.object(
            external_handler, "close", wraps=external_handler.close
        ) as close:
            self.manager.update_config({"backup_count": 3})
            close.assert_not_called()

        self.assertNotIn(external_handler, self.manager.logger.handlers)
        self.assertFalse(external_handler.stream.closed)
        other_logger.info("external sink still usable")
        self.assertIn("external sink still usable", external_path.read_text())

    def test_level_only_change_keeps_current_file_handler(self) -> None:
        handler = next(
            h
            for h in self.manager.logger.handlers
            if isinstance(h, logging.FileHandler)
        )
        stream = handler.stream

        self.manager.update_config({"level": "DEBUG"})

        self.assertIn(handler, self.manager.logger.handlers)
        self.assertIs(handler.stream, stream)
        self.assertFalse(stream.closed)
        self.logger.debug("same sink still usable")
        self.assertIn(
            "same sink still usable",
            pathlib.Path(self.log_dir, "prompt_manager.log").read_text(),
        )

    def test_replacement_manager_closes_previous_owned_handlers(self) -> None:
        previous = self.manager.logger.handlers[:]
        for handler in previous:
            self.addCleanup(handler.close)
        old_file = next(h for h in previous if isinstance(h, logging.FileHandler))
        stream = old_file.stream
        self.assertFalse(stream.closed)
        fake_file = pathlib.Path(self.tmp.name, "pkg", "utils", "logging_config.py")

        with mock.patch.object(logging_config, "__file__", str(fake_file)):
            replacement = object.__new__(PromptManagerLogger)
            replacement.__init__()

        self.assertIs(replacement.logger, self.manager.logger)
        self.assertTrue(stream.closed)
        self.assertIsNone(old_file.stream)
        self.assertTrue(all(h._closed for h in previous))
        self._filter_to_this_thread()
        self.logger.info("replacement manager sink")
        self.assertEqual(
            replacement.get_recent_logs(limit=1)[0]["message"],
            "replacement manager sink",
        )

    def test_level_change_applies_to_every_prompt_manager_logger(self):
        self.manager.update_config({"level": "DEBUG"})
        self.assertEqual(logging.getLogger("prompt_manager").level, logging.DEBUG)
        self.assertEqual(
            logging.getLogger("prompt_manager.database").level, logging.DEBUG
        )
        self.assertEqual(self.manager.get_config()["level"], "DEBUG")
        self.logger.debug("now visible")
        self.assertEqual(self.manager.get_recent_logs(limit=1)[0]["level"], "DEBUG")
        self.manager.update_config({"level": "bogus"})  # falls back to INFO
        self.assertEqual(logging.getLogger("prompt_manager").level, logging.INFO)

    def test_get_config_returns_a_copy(self):
        config = self.manager.get_config()
        config["level"] = "CRITICAL"
        self.assertEqual(self.manager.get_config()["level"], "INFO")

    def test_handler_flags_rebuild_the_handlers(self):
        root = logging.getLogger("prompt_manager")
        self.assertFalse(any(type(h) is logging.StreamHandler for h in root.handlers))
        self.assertTrue(any(isinstance(h, logging.FileHandler) for h in root.handlers))
        self.manager.update_config({"console_logging": True, "file_logging": False})
        self.assertTrue(any(type(h) is logging.StreamHandler for h in root.handlers))
        self.assertFalse(any(isinstance(h, logging.FileHandler) for h in root.handlers))
        self.assertTrue(any(isinstance(h, MemoryBufferHandler) for h in root.handlers))
        self.manager.update_config({"console_logging": False})


class TestConsoleNoise(IsolatedLoggerTestCase):
    """ComfyUI's console only sees warnings unless the level is DEBUG."""

    def _console_handler(self) -> logging.StreamHandler:
        return next(
            h
            for h in logging.getLogger("prompt_manager").handlers
            if type(h) is logging.StreamHandler
        )

    def test_records_do_not_propagate_to_comfyui_root_logger(self) -> None:
        self.assertFalse(logging.getLogger("prompt_manager").propagate)

    def test_console_hides_info_at_the_default_level(self) -> None:
        self.manager.update_config({"console_logging": True})
        self.assertEqual(self._console_handler().level, logging.WARNING)
        self.logger.info("startup chatter")
        self.assertEqual(self.manager.get_recent_logs(limit=1)[0]["level"], "INFO")
        log_file = pathlib.Path(self.log_dir) / "prompt_manager.log"
        self.assertIn("startup chatter", log_file.read_text())
        self.manager.update_config({"console_logging": False})

    def test_debug_level_opens_the_console(self) -> None:
        self.manager.update_config({"console_logging": True, "level": "DEBUG"})
        self.assertEqual(self._console_handler().level, logging.DEBUG)
        self.manager.update_config({"level": "INFO"})
        self.assertEqual(self._console_handler().level, logging.WARNING)
        self.manager.update_config({"console_logging": False})


if __name__ == "__main__":
    unittest.main()