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, }, )