1  from pathlib import Path
  2  from unittest.mock import Mock, patch
  3  
  4  import pytest
  5  
  6  from mlflow.telemetry.constant import (
  7      CONFIG_STAGING_URL,
  8      CONFIG_URL,
  9      FALLBACK_UI_CONFIG,
 10      UI_CONFIG_STAGING_URL,
 11      UI_CONFIG_URL,
 12  )
 13  from mlflow.telemetry.schemas import Environment
 14  from mlflow.telemetry.utils import (
 15      _detect_environment,
 16      _get_config_url,
 17      fetch_ui_telemetry_config,
 18      is_telemetry_disabled,
 19  )
 20  
 21  
 22  @pytest.mark.parametrize(
 23      ("env_var", "expected"),
 24      [
 25          ("KAGGLE_KERNEL_RUN_TYPE", Environment.KAGGLE.value),
 26          ("COLAB_RELEASE_TAG", Environment.COLAB.value),
 27          ("AZUREML_FRAMEWORK", Environment.AZURE_ML.value),
 28          ("SAGEMAKER_APP_TYPE", Environment.SAGEMAKER_STUDIO.value),
 29      ],
 30  )
 31  def test_detect_environment_from_env_var(monkeypatch, env_var, expected):
 32      monkeypatch.setenv(env_var, "some_value")
 33      assert _detect_environment() == expected
 34  
 35  
 36  def test_detect_environment_sagemaker_notebook(tmp_path, monkeypatch):
 37      metadata_file = tmp_path / "resource-metadata.json"
 38      metadata_file.write_text("{}")
 39      with patch("mlflow.telemetry.utils.Path") as mock_path:
 40  
 41          def path_side_effect(p):
 42              if p == "/opt/ml/metadata/resource-metadata.json":
 43                  return metadata_file
 44              return Path(p)
 45  
 46          mock_path.side_effect = path_side_effect
 47          assert _detect_environment() == Environment.SAGEMAKER_NOTEBOOK
 48  
 49  
 50  def test_detect_environment_docker(tmp_path, monkeypatch):
 51      dockerenv = tmp_path / ".dockerenv"
 52      dockerenv.touch()
 53      with patch("mlflow.telemetry.utils.Path") as mock_path:
 54  
 55          def path_side_effect(p):
 56              if p == "/.dockerenv":
 57                  return dockerenv
 58              return Path(p)
 59  
 60          mock_path.side_effect = path_side_effect
 61          assert _detect_environment() == Environment.DOCKER
 62  
 63  
 64  def test_detect_environment_none():
 65      assert _detect_environment() is None
 66  
 67  
 68  def test_is_telemetry_disabled(monkeypatch, bypass_env_check):
 69      assert is_telemetry_disabled() is False
 70  
 71      with monkeypatch.context() as m:
 72          m.setenv("MLFLOW_DISABLE_TELEMETRY", "true")
 73          assert is_telemetry_disabled() is True
 74  
 75      assert is_telemetry_disabled() is False
 76  
 77      with monkeypatch.context() as m:
 78          m.setenv("DO_NOT_TRACK", "true")
 79          assert is_telemetry_disabled() is True
 80  
 81  
 82  def test_get_config_url(bypass_env_check):
 83      assert _get_config_url("1.0.0") == f"{CONFIG_URL}/1.0.0.json"
 84      assert _get_config_url("1.0.0.rc0") == f"{CONFIG_URL}/1.0.0.rc0.json"
 85      assert _get_config_url("1.0.0.dev0") == f"{CONFIG_STAGING_URL}/1.0.0.dev0.json"
 86      assert _get_config_url("1.0.0+abc") is None
 87      assert _get_config_url("1.0.0", is_ui=True) == f"{UI_CONFIG_URL}/1.0.0.json"
 88      assert _get_config_url("1.0.0.rc0", is_ui=True) == f"{UI_CONFIG_URL}/1.0.0.rc0.json"
 89      assert _get_config_url("1.0.0.dev0", is_ui=True) == f"{UI_CONFIG_STAGING_URL}/1.0.0.dev0.json"
 90      assert _get_config_url("1.0.0+abc", is_ui=True) is None
 91  
 92  
 93  def test_fetch_ui_telemetry_config_fetch_success(bypass_env_check):
 94      mock_config = {
 95          "mlflow_version": "3.7.1.dev0",
 96          "disable_telemetry": False,
 97          "rollout_percentage": 100,
 98          "ingestion_url": "https://api.mlflow-telemetry.io/staging/log",
 99          "disable_sdks": [],
100          "disable_os": [],
101          "disable_events": [],
102          "disable_ui_telemetry": False,
103          "disable_ui_events": ["test_event"],
104          "ui_rollout_percentage": 100,
105      }
106  
107      with patch("requests.get") as mock_get:
108          mock_response = Mock()
109          mock_response.status_code = 200
110          mock_response.json.return_value = mock_config
111          mock_get.return_value = mock_response
112  
113          result = fetch_ui_telemetry_config()
114  
115          assert result["disable_ui_telemetry"] is False
116          assert result["disable_ui_events"] == ["test_event"]
117          assert result["ui_rollout_percentage"] == 100
118  
119  
120  def test_fetch_ui_telemetry_config_fetch_error_fallback(bypass_env_check):
121      with patch("requests.get") as mock_get:
122          mock_response = Mock()
123          mock_response.status_code = 404
124          mock_get.return_value = mock_response
125  
126          result = fetch_ui_telemetry_config()
127  
128          assert result == FALLBACK_UI_CONFIG
129          assert result["disable_ui_telemetry"] is True
130  
131      with patch("requests.get") as mock_get:
132          mock_get.side_effect = Exception("Network error")
133  
134          result = fetch_ui_telemetry_config()
135  
136          assert result == FALLBACK_UI_CONFIG
137          assert result["disable_ui_telemetry"] is True
138          assert result["disable_ui_events"] == []
139          assert result["ui_rollout_percentage"] == 0