mlflow--mlflow/tests/utils/test_logging_utils.py

import logging
import os
import re
import subprocess
import sys
import uuid
from io import StringIO

import pytest

import mlflow
from mlflow.utils import logging_utils
from mlflow.utils.logging_utils import (
    LOGGING_LINE_FORMAT,
    SensitiveQueryParamFilter,
    _redact_sensitive_query_params,
    eprint,
    suppress_logs,
)

logger = logging.getLogger(mlflow.__name__)

LOGGING_FNS_TO_TEST = [logger.info, logger.warning, logger.critical, eprint]


@pytest.fixture(autouse=True)
def reset_stderr():
    prev_stderr = sys.stderr
    yield
    sys.stderr = prev_stderr


@pytest.fixture(autouse=True)
def reset_logging_enablement():
    yield
    logging_utils.enable_logging()


@pytest.fixture(autouse=True)
def reset_logging_level():
    level_before = logger.level
    yield
    logger.setLevel(level_before)


class SampleStream:
    def __init__(self):
        self.content = None
        self.flush_count = 0

    def write(self, text):
        self.content = (self.content or "") + text

    def flush(self):
        self.flush_count += 1

    def reset(self):
        self.content = None
        self.flush_count = 0


@pytest.mark.parametrize("logging_fn", LOGGING_FNS_TO_TEST)
def test_event_logging_apis_respect_stderr_reassignment(logging_fn):
    stream1 = SampleStream()
    stream2 = SampleStream()
    message_content = "test message"

    sys.stderr = stream1
    assert stream1.content is None
    logging_fn(message_content)
    assert message_content in stream1.content
    assert stream2.content is None
    stream1.reset()

    sys.stderr = stream2
    assert stream2.content is None
    logging_fn(message_content)
    assert message_content in stream2.content
    assert stream1.content is None


@pytest.mark.parametrize("logging_fn", LOGGING_FNS_TO_TEST)
def test_event_logging_apis_respect_stream_disablement_enablement(logging_fn):
    stream = SampleStream()
    sys.stderr = stream
    message_content = "test message"

    assert stream.content is None
    logging_fn(message_content)
    assert message_content in stream.content
    stream.reset()

    logging_utils.disable_logging()
    logging_fn(message_content)
    assert stream.content is None
    stream.reset()

    logging_utils.enable_logging()
    assert stream.content is None
    logging_fn(message_content)
    assert message_content in stream.content


def test_event_logging_stream_flushes_properly():
    stream = SampleStream()
    sys.stderr = stream

    eprint("foo", flush=True)
    assert "foo" in stream.content
    assert stream.flush_count > 0


def test_debug_logs_emitted_correctly_when_configured():
    stream = SampleStream()
    sys.stderr = stream

    logger.setLevel(logging.DEBUG)
    logger.debug("test debug")
    assert "test debug" in stream.content


def test_suppress_logs():
    module = "test_logger"
    logger = logging.getLogger(module)

    message = "This message should be suppressed."

    capture_stream = StringIO()
    stream_handler = logging.StreamHandler(capture_stream)
    logger.addHandler(stream_handler)

    logger.error(message)
    assert message in capture_stream.getvalue()

    capture_stream.truncate(0)
    with suppress_logs(module, re.compile(r"This .* be suppressed.")):
        logger.error(message)
    assert len(capture_stream.getvalue()) == 0


@pytest.mark.parametrize(
    ("log_level", "expected"),
    [
        ("DEBUG", True),
        ("INFO", False),
        ("NOTSET", False),
    ],
)
def test_logging_level(log_level: str, expected: bool) -> None:
    random_str = str(uuid.uuid4())
    stdout = subprocess.check_output(
        [
            sys.executable,
            "-c",
            f"from mlflow.utils.logging_utils import _debug; _debug({random_str!r})",
        ],
        env=os.environ.copy() | {"MLFLOW_LOGGING_LEVEL": log_level},
        stderr=subprocess.STDOUT,
        text=True,
    )

    assert (random_str in stdout) is expected


@pytest.mark.parametrize(
    "env_var_name",
    ["MLFLOW_CONFIGURE_LOGGING", "MLFLOW_LOGGING_CONFIGURE_LOGGING"],
)
@pytest.mark.parametrize(
    "value",
    ["0", "1"],
)
def test_mlflow_configure_logging_env_var(env_var_name: str, value: str) -> None:
    expected_level = logging.INFO if value == "1" else logging.WARNING
    subprocess.check_call(
        [
            sys.executable,
            "-c",
            f"""
import logging
import mlflow

assert logging.getLogger("mlflow").isEnabledFor({expected_level})
""",
        ],
        env=os.environ.copy() | {env_var_name: value},
    )


@pytest.mark.parametrize(
    ("message", "expected"),
    [
        (
            "Retrying after broken: /path?X-Amz-Security-Token=ABC&"
            "X-Amz-Credential=KEY%2F20260414%2Fap-northeast-1%2Fs3%2Faws4_request&"
            "X-Amz-Date=20260414T154032Z&X-Amz-Signature=DEADBEEF",
            "Retrying after broken: /path?X-Amz-Security-Token=[REDACTED]&"
            "X-Amz-Credential=[REDACTED]&"
            "X-Amz-Date=20260414T154032Z&X-Amz-Signature=[REDACTED]",
        ),
        (
            "https://store.googleapis.com/o?X-Goog-Signature=ZZZ&X-Goog-Credential=YYY",
            "https://store.googleapis.com/o?X-Goog-Signature=[REDACTED]&"
            "X-Goog-Credential=[REDACTED]",
        ),
        (
            "https://store.blob.core.windows.net/c/b?sig=AAA&se=2026-04-14T16%3A40%3A32Z",
            "https://store.blob.core.windows.net/c/b?sig=[REDACTED]&se=2026-04-14T16%3A40%3A32Z",
        ),
        (
            "just a normal log message without credentials",
            "just a normal log message without credentials",
        ),
        ("path?foo=bar&baz=qux", "path?foo=bar&baz=qux"),
    ],
)
def test_redact_sensitive_query_params(message: str, expected: str):
    assert _redact_sensitive_query_params(message) == expected


def test_sensitive_query_param_filter_rewrites_record_message():
    capture_stream = StringIO()
    logger = logging.getLogger(f"test_urllib3_{uuid.uuid4().hex}")
    logger.addHandler(logging.StreamHandler(capture_stream))
    logger.addFilter(SensitiveQueryParamFilter())
    logger.setLevel(logging.WARNING)

    logger.warning(
        "Retrying (%r) after connection broken by %r: %s",
        "Retry(total=1)",
        "NameResolutionError(...)",
        "/tokyo-prod/artifacts/traces.json"
        "?X-Amz-Security-Token=SECRET_TOKEN"
        "&X-Amz-Algorithm=AWS4-HMAC-SHA256"
        "&X-Amz-Signature=SECRET_SIGNATURE"
        "&X-Amz-Credential=SECRET_CREDENTIAL",
    )

    output = capture_stream.getvalue()
    assert "X-Amz-Security-Token=[REDACTED]" in output
    assert "X-Amz-Signature=[REDACTED]" in output
    assert "X-Amz-Credential=[REDACTED]" in output
    assert "X-Amz-Algorithm=AWS4-HMAC-SHA256" in output
    assert "SECRET_TOKEN" not in output
    assert "SECRET_SIGNATURE" not in output
    assert "SECRET_CREDENTIAL" not in output


def test_sensitive_query_param_filter_installed_on_urllib3_connectionpool():
    urllib3_logger = logging.getLogger("urllib3.connectionpool")
    assert any(isinstance(f, SensitiveQueryParamFilter) for f in urllib3_logger.filters)


@pytest.mark.parametrize("configure_logging", ["0", "1"])
def test_alembic_logging_respects_configure_flag(configure_logging: str, tmp_sqlite_uri: str):
    user_specified_format = "CUSTOM: %(name)s - %(message)s"
    actual_format = user_specified_format if configure_logging == "0" else LOGGING_LINE_FORMAT
    code = f"""
import logging

# user-specified format, this should only take effect if configure_logging is 0
logging.basicConfig(level=logging.INFO, format={user_specified_format!r})

import mlflow

# Check the alembic logger format, which is now configured in _configure_mlflow_loggers
alembic_logger = logging.getLogger("alembic")
if {configure_logging!r} == "1":
    # When MLFLOW_CONFIGURE_LOGGING is enabled, alembic logger has its own handler
    assert len(alembic_logger.handlers) > 0
    actual_format = alembic_logger.handlers[0].formatter._fmt
else:
    # When MLFLOW_CONFIGURE_LOGGING is disabled, alembic logger propagates to root
    assert alembic_logger.propagate
    root_logger = logging.getLogger()
    actual_format = root_logger.handlers[0].formatter._fmt

assert actual_format == {actual_format!r}, actual_format
"""
    subprocess.check_call(
        [sys.executable, "-c", code],
        env={
            **os.environ,
            "MLFLOW_TRACKING_URI": tmp_sqlite_uri,
            "MLFLOW_CONFIGURE_LOGGING": configure_logging,
        },
    )