mlflow--mlflow/tests/telemetry/test_utils.py

from pathlib import Path
from unittest.mock import Mock, patch

import pytest

import mlflow.telemetry.utils
from mlflow.telemetry.constant import (
    CONFIG_STAGING_URL,
    CONFIG_URL,
    FALLBACK_UI_CONFIG,
    UI_CONFIG_STAGING_URL,
    UI_CONFIG_URL,
)
from mlflow.telemetry.schemas import Environment
from mlflow.telemetry.utils import (
    _detect_environment,
    _get_config_url,
    fetch_ui_telemetry_config,
    is_telemetry_disabled,
)


@pytest.mark.parametrize(
    ("env_var", "expected"),
    [
        ("KAGGLE_KERNEL_RUN_TYPE", Environment.KAGGLE.value),
        ("COLAB_RELEASE_TAG", Environment.COLAB.value),
        ("AZUREML_FRAMEWORK", Environment.AZURE_ML.value),
        ("SAGEMAKER_APP_TYPE", Environment.SAGEMAKER_STUDIO.value),
    ],
)
def test_detect_environment_from_env_var(monkeypatch, env_var, expected):
    monkeypatch.setenv(env_var, "some_value")
    assert _detect_environment() == expected


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        ("demo", Environment.DEMO.value),
        ("staging", None),
    ],
)
def test_detect_environment_deployment_env(monkeypatch, value, expected):
    monkeypatch.setenv("MLFLOW_DEPLOYMENT_ENV", value)
    assert _detect_environment() == expected


def test_detect_environment_sagemaker_notebook(tmp_path, monkeypatch):
    metadata_file = tmp_path / "resource-metadata.json"
    metadata_file.write_text("{}")
    with patch("mlflow.telemetry.utils.Path") as mock_path:

        def path_side_effect(p):
            if p == "/opt/ml/metadata/resource-metadata.json":
                return metadata_file
            return Path(p)

        mock_path.side_effect = path_side_effect
        assert _detect_environment() == Environment.SAGEMAKER_NOTEBOOK


def test_detect_environment_docker(tmp_path, monkeypatch):
    dockerenv = tmp_path / ".dockerenv"
    dockerenv.touch()
    with patch("mlflow.telemetry.utils.Path") as mock_path:

        def path_side_effect(p):
            if p == "/.dockerenv":
                return dockerenv
            return Path(p)

        mock_path.side_effect = path_side_effect
        assert _detect_environment() == Environment.DOCKER


def test_detect_environment_none():
    assert _detect_environment() is None


def test_is_telemetry_disabled(monkeypatch, bypass_env_check):
    assert is_telemetry_disabled() is False

    with monkeypatch.context() as m:
        m.setenv("MLFLOW_DISABLE_TELEMETRY", "true")
        assert is_telemetry_disabled() is True

    assert is_telemetry_disabled() is False

    with monkeypatch.context() as m:
        m.setenv("DO_NOT_TRACK", "true")
        assert is_telemetry_disabled() is True


def test_is_telemetry_disabled_does_not_disable_in_databricks(monkeypatch, bypass_env_check):
    # Running inside DBR / model serving must not disable telemetry — events
    # are routed to the workspace via _forward_to_databricks. The OSS branch
    # is independently gated inside TelemetryClient._process_records.
    monkeypatch.setattr(mlflow.telemetry.utils, "_IS_IN_DATABRICKS", True)
    assert is_telemetry_disabled() is False


def test_get_config_url(bypass_env_check):
    assert _get_config_url("1.0.0") == f"{CONFIG_URL}/1.0.0.json"
    assert _get_config_url("1.0.0.rc0") == f"{CONFIG_URL}/1.0.0.rc0.json"
    assert _get_config_url("1.0.0.dev0") == f"{CONFIG_STAGING_URL}/1.0.0.dev0.json"
    assert _get_config_url("1.0.0+abc") is None
    assert _get_config_url("1.0.0", is_ui=True) == f"{UI_CONFIG_URL}/1.0.0.json"
    assert _get_config_url("1.0.0.rc0", is_ui=True) == f"{UI_CONFIG_URL}/1.0.0.rc0.json"
    assert _get_config_url("1.0.0.dev0", is_ui=True) == f"{UI_CONFIG_STAGING_URL}/1.0.0.dev0.json"
    assert _get_config_url("1.0.0+abc", is_ui=True) is None


def test_fetch_ui_telemetry_config_fetch_success(bypass_env_check):
    mock_config = {
        "mlflow_version": "3.7.1.dev0",
        "disable_telemetry": False,
        "rollout_percentage": 100,
        "ingestion_url": "https://api.mlflow-telemetry.io/staging/log",
        "disable_sdks": [],
        "disable_os": [],
        "disable_events": [],
        "disable_ui_telemetry": False,
        "disable_ui_events": ["test_event"],
        "ui_rollout_percentage": 100,
    }

    with patch("requests.get") as mock_get:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = mock_config
        mock_get.return_value = mock_response

        result = fetch_ui_telemetry_config()

        assert result["disable_ui_telemetry"] is False
        assert result["disable_ui_events"] == ["test_event"]
        assert result["ui_rollout_percentage"] == 100


def test_fetch_ui_telemetry_config_fetch_error_fallback(bypass_env_check):
    with patch("requests.get") as mock_get:
        mock_response = Mock()
        mock_response.status_code = 404
        mock_get.return_value = mock_response

        result = fetch_ui_telemetry_config()

        assert result == FALLBACK_UI_CONFIG
        assert result["disable_ui_telemetry"] is True

    with patch("requests.get") as mock_get:
        mock_get.side_effect = Exception("Network error")

        result = fetch_ui_telemetry_config()

        assert result == FALLBACK_UI_CONFIG
        assert result["disable_ui_telemetry"] is True
        assert result["disable_ui_events"] == []
        assert result["ui_rollout_percentage"] == 0