1  import json
  2  
  3  import pytest
  4  
  5  from mlflow.claude_code.config import (
  6      MLFLOW_TRACING_ENABLED,
  7      get_env_var,
  8      get_tracing_status,
  9      load_claude_config,
 10      save_claude_config,
 11      setup_environment_config,
 12  )
 13  
 14  
 15  @pytest.fixture
 16  def temp_settings_path(tmp_path):
 17      """Provide a temporary settings.json path for tests."""
 18      return tmp_path / "settings.json"
 19  
 20  
 21  def test_load_claude_config_valid_json(temp_settings_path):
 22      config_data = {"tools": {"computer_20241022": {"name": "computer"}}}
 23      with open(temp_settings_path, "w") as f:
 24          json.dump(config_data, f)
 25  
 26      result = load_claude_config(temp_settings_path)
 27      assert result == config_data
 28  
 29  
 30  def test_load_claude_config_missing_file(tmp_path):
 31      non_existent_path = tmp_path / "non_existent.json"
 32      result = load_claude_config(non_existent_path)
 33      assert result == {}
 34  
 35  
 36  def test_load_claude_config_invalid_json(temp_settings_path):
 37      with open(temp_settings_path, "w") as f:
 38          f.write("invalid json content")
 39  
 40      result = load_claude_config(temp_settings_path)
 41      assert result == {}
 42  
 43  
 44  def test_save_claude_config_creates_file(temp_settings_path):
 45      config_data = {"test": "value"}
 46      save_claude_config(temp_settings_path, config_data)
 47  
 48      assert temp_settings_path.exists()
 49      saved_data = json.loads(temp_settings_path.read_text())
 50      assert saved_data == config_data
 51  
 52  
 53  def test_save_claude_config_creates_directory(tmp_path):
 54      nested_path = tmp_path / "nested" / "dir" / "settings.json"
 55      config_data = {"test": "value"}
 56  
 57      save_claude_config(nested_path, config_data)
 58  
 59      assert nested_path.exists()
 60      saved_data = json.loads(nested_path.read_text())
 61      assert saved_data == config_data
 62  
 63  
 64  def test_get_env_var_from_os_environment_when_no_settings(tmp_path, monkeypatch):
 65      monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "test_os_value")
 66      monkeypatch.chdir(tmp_path)
 67  
 68      result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
 69      assert result == "test_os_value"
 70  
 71  
 72  def test_get_env_var_settings_takes_precedence_over_os_env(tmp_path, monkeypatch):
 73      monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")
 74  
 75      config_data = {"env": {MLFLOW_TRACING_ENABLED: "settings_value"}}
 76      claude_settings_path = tmp_path / ".claude" / "settings.json"
 77      claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
 78      with open(claude_settings_path, "w") as f:
 79          json.dump(config_data, f)
 80  
 81      monkeypatch.chdir(tmp_path)
 82      result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
 83      assert result == "settings_value"
 84  
 85  
 86  def test_get_env_var_falls_back_to_os_env_when_not_in_settings(tmp_path, monkeypatch):
 87      monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")
 88  
 89      config_data = {"env": {"OTHER_VAR": "other_value"}}
 90      claude_settings_path = tmp_path / ".claude" / "settings.json"
 91      claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
 92      with open(claude_settings_path, "w") as f:
 93          json.dump(config_data, f)
 94  
 95      monkeypatch.chdir(tmp_path)
 96      result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
 97      assert result == "os_value"
 98  
 99  
100  def test_get_env_var_default_when_not_found(tmp_path, monkeypatch):
101      # Ensure OS env var is not set
102      monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)
103  
104      # Create empty settings file in .claude directory
105      claude_settings_path = tmp_path / ".claude" / "settings.json"
106      claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
107      with open(claude_settings_path, "w") as f:
108          json.dump({}, f)
109  
110      # Change to temp directory so .claude/settings.json is found
111      monkeypatch.chdir(tmp_path)
112      result = get_env_var(MLFLOW_TRACING_ENABLED, "default_value")
113      assert result == "default_value"
114  
115  
116  def test_get_tracing_status_enabled(temp_settings_path):
117      # Create settings with tracing enabled
118      config_data = {"env": {MLFLOW_TRACING_ENABLED: "true"}}
119      with open(temp_settings_path, "w") as f:
120          json.dump(config_data, f)
121  
122      status = get_tracing_status(temp_settings_path)
123      assert status.enabled is True
124      assert hasattr(status, "tracking_uri")
125  
126  
127  def test_get_tracing_status_disabled(temp_settings_path):
128      # Create settings with tracing disabled
129      config_data = {"env": {MLFLOW_TRACING_ENABLED: "false"}}
130      with open(temp_settings_path, "w") as f:
131          json.dump(config_data, f)
132  
133      status = get_tracing_status(temp_settings_path)
134      assert status.enabled is False
135  
136  
137  def test_get_tracing_status_no_config(tmp_path):
138      non_existent_path = tmp_path / "missing.json"
139      status = get_tracing_status(non_existent_path)
140      assert status.enabled is False
141      assert status.reason == "No configuration found"
142  
143  
144  def test_setup_environment_config_new_file(temp_settings_path):
145      tracking_uri = "test://localhost"
146      experiment_id = "123"
147  
148      setup_environment_config(temp_settings_path, tracking_uri, experiment_id)
149  
150      # Verify file was created
151      assert temp_settings_path.exists()
152  
153      # Verify configuration contents
154      config = json.loads(temp_settings_path.read_text())
155  
156      env_vars = config["env"]
157      assert env_vars[MLFLOW_TRACING_ENABLED] == "true"
158      assert env_vars["MLFLOW_TRACKING_URI"] == tracking_uri
159      assert env_vars["MLFLOW_EXPERIMENT_ID"] == experiment_id
160  
161  
162  def test_setup_environment_config_experiment_id_precedence(temp_settings_path):
163      # Create existing config with different experiment ID
164      existing_config = {
165          "env": {
166              MLFLOW_TRACING_ENABLED: "true",
167              "MLFLOW_EXPERIMENT_ID": "old_id",
168              "MLFLOW_TRACKING_URI": "old_uri",
169          }
170      }
171      with open(temp_settings_path, "w") as f:
172          json.dump(existing_config, f)
173  
174      new_tracking_uri = "new://localhost"
175      new_experiment_id = "new_id"
176  
177      setup_environment_config(temp_settings_path, new_tracking_uri, new_experiment_id)
178  
179      # Verify configuration was updated
180      config = json.loads(temp_settings_path.read_text())
181  
182      env_vars = config["env"]
183      assert env_vars[MLFLOW_TRACING_ENABLED] == "true"
184      assert env_vars["MLFLOW_TRACKING_URI"] == new_tracking_uri
185      assert env_vars["MLFLOW_EXPERIMENT_ID"] == new_experiment_id