mlflow--mlflow/tests/claude_code/test_config.py

import json

import pytest

from mlflow.claude_code.config import (
    MLFLOW_TRACING_ENABLED,
    get_env_var,
    get_tracing_status,
    load_claude_config,
    save_claude_config,
    setup_environment_config,
)


@pytest.fixture
def temp_settings_path(tmp_path):
    """Provide a temporary settings.json path for tests."""
    return tmp_path / "settings.json"


def test_load_claude_config_valid_json(temp_settings_path):
    config_data = {"tools": {"computer_20241022": {"name": "computer"}}}
    with open(temp_settings_path, "w") as f:
        json.dump(config_data, f)

    result = load_claude_config(temp_settings_path)
    assert result == config_data


def test_load_claude_config_missing_file(tmp_path):
    non_existent_path = tmp_path / "non_existent.json"
    result = load_claude_config(non_existent_path)
    assert result == {}


def test_load_claude_config_invalid_json(temp_settings_path):
    with open(temp_settings_path, "w") as f:
        f.write("invalid json content")

    result = load_claude_config(temp_settings_path)
    assert result == {}


def test_save_claude_config_creates_file(temp_settings_path):
    config_data = {"test": "value"}
    save_claude_config(temp_settings_path, config_data)

    assert temp_settings_path.exists()
    saved_data = json.loads(temp_settings_path.read_text())
    assert saved_data == config_data


def test_save_claude_config_creates_directory(tmp_path):
    nested_path = tmp_path / "nested" / "dir" / "settings.json"
    config_data = {"test": "value"}

    save_claude_config(nested_path, config_data)

    assert nested_path.exists()
    saved_data = json.loads(nested_path.read_text())
    assert saved_data == config_data


def test_get_env_var_from_os_environment_when_no_settings(tmp_path, monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "test_os_value")
    monkeypatch.chdir(tmp_path)

    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
    assert result == "test_os_value"


def test_get_env_var_os_env_takes_precedence_over_settings(tmp_path, monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")

    config_data = {"env": {MLFLOW_TRACING_ENABLED: "settings_value"}}
    claude_settings_path = tmp_path / ".claude" / "settings.json"
    claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
    with open(claude_settings_path, "w") as f:
        json.dump(config_data, f)

    monkeypatch.chdir(tmp_path)
    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
    assert result == "os_value"


def test_get_env_var_falls_back_to_os_env_when_not_in_settings(tmp_path, monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")

    config_data = {"env": {"OTHER_VAR": "other_value"}}
    claude_settings_path = tmp_path / ".claude" / "settings.json"
    claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
    with open(claude_settings_path, "w") as f:
        json.dump(config_data, f)

    monkeypatch.chdir(tmp_path)
    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
    assert result == "os_value"


def test_get_env_var_from_settings_local_json(tmp_path, monkeypatch):
    monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)

    config_data = {"env": {MLFLOW_TRACING_ENABLED: "local_value"}}
    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)
    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump(config_data, f)

    monkeypatch.chdir(tmp_path)
    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
    assert result == "local_value"


def test_get_env_var_settings_local_overrides_settings(tmp_path, monkeypatch):
    monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)

    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)

    with open(claude_dir / "settings.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "shared_value"}}, f)
    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "local_value"}}, f)

    monkeypatch.chdir(tmp_path)
    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
    assert result == "local_value"


def test_get_env_var_falls_through_local_to_settings(tmp_path, monkeypatch):
    monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False)

    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)

    # local has ENABLED, shared has URI
    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)
    with open(claude_dir / "settings.json", "w") as f:
        json.dump({"env": {"MLFLOW_TRACKING_URI": "https://example.com"}}, f)

    monkeypatch.chdir(tmp_path)
    assert get_env_var(MLFLOW_TRACING_ENABLED, "default") == "true"
    assert get_env_var("MLFLOW_TRACKING_URI", "default") == "https://example.com"


def test_get_env_var_os_env_takes_precedence_over_local(tmp_path, monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")

    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)
    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "local_value"}}, f)

    monkeypatch.chdir(tmp_path)
    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
    assert result == "os_value"


def test_get_tracing_status_merges_local_and_shared(tmp_path):
    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)
    settings_path = claude_dir / "settings.json"

    # URI in shared, ENABLED in local
    with open(settings_path, "w") as f:
        json.dump({"env": {"MLFLOW_TRACKING_URI": "https://example.com"}}, f)
    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)

    status = get_tracing_status(settings_path)
    assert status.enabled is True
    assert status.tracking_uri == "https://example.com"


def test_get_tracing_status_local_overrides_shared(tmp_path):
    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)
    settings_path = claude_dir / "settings.json"

    with open(settings_path, "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)
    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "false"}}, f)

    status = get_tracing_status(settings_path)
    assert status.enabled is False


def test_get_tracing_status_from_local_only(tmp_path):
    claude_dir = tmp_path / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)
    settings_path = claude_dir / "settings.json"

    with open(claude_dir / "settings.local.json", "w") as f:
        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)

    status = get_tracing_status(settings_path)
    assert status.enabled is True


def test_get_env_var_default_when_not_found(tmp_path, monkeypatch):
    # Ensure OS env var is not set
    monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)

    # Create empty settings file in .claude directory
    claude_settings_path = tmp_path / ".claude" / "settings.json"
    claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
    with open(claude_settings_path, "w") as f:
        json.dump({}, f)

    # Change to temp directory so .claude/settings.json is found
    monkeypatch.chdir(tmp_path)
    result = get_env_var(MLFLOW_TRACING_ENABLED, "default_value")
    assert result == "default_value"


def test_get_tracing_status_enabled(temp_settings_path):
    # Create settings with tracing enabled
    config_data = {"env": {MLFLOW_TRACING_ENABLED: "true"}}
    with open(temp_settings_path, "w") as f:
        json.dump(config_data, f)

    status = get_tracing_status(temp_settings_path)
    assert status.enabled is True
    assert hasattr(status, "tracking_uri")


def test_get_tracing_status_disabled(temp_settings_path):
    # Create settings with tracing disabled
    config_data = {"env": {MLFLOW_TRACING_ENABLED: "false"}}
    with open(temp_settings_path, "w") as f:
        json.dump(config_data, f)

    status = get_tracing_status(temp_settings_path)
    assert status.enabled is False


def test_get_tracing_status_no_config(tmp_path):
    non_existent_path = tmp_path / "missing.json"
    status = get_tracing_status(non_existent_path)
    assert status.enabled is False
    assert status.reason == "No configuration found"


def test_setup_environment_config_new_file(temp_settings_path):
    tracking_uri = "test://localhost"
    experiment_id = "123"

    setup_environment_config(temp_settings_path, tracking_uri, experiment_id)

    # Verify file was created
    assert temp_settings_path.exists()

    # Verify configuration contents
    config = json.loads(temp_settings_path.read_text())

    env_vars = config["env"]
    assert env_vars[MLFLOW_TRACING_ENABLED] == "true"
    assert env_vars["MLFLOW_TRACKING_URI"] == tracking_uri
    assert env_vars["MLFLOW_EXPERIMENT_ID"] == experiment_id


def test_setup_environment_config_experiment_id_precedence(temp_settings_path):
    # Create existing config with different experiment ID
    existing_config = {
        "env": {
            MLFLOW_TRACING_ENABLED: "true",
            "MLFLOW_EXPERIMENT_ID": "old_id",
            "MLFLOW_TRACKING_URI": "old_uri",
        }
    }
    with open(temp_settings_path, "w") as f:
        json.dump(existing_config, f)

    new_tracking_uri = "new://localhost"
    new_experiment_id = "new_id"

    setup_environment_config(temp_settings_path, new_tracking_uri, new_experiment_id)

    # Verify configuration was updated
    config = json.loads(temp_settings_path.read_text())

    env_vars = config["env"]
    assert env_vars[MLFLOW_TRACING_ENABLED] == "true"
    assert env_vars["MLFLOW_TRACKING_URI"] == new_tracking_uri
    assert env_vars["MLFLOW_EXPERIMENT_ID"] == new_experiment_id


def test_setup_environment_config_defaults_to_current_tracking_uri_and_default_experiment(
    temp_settings_path, monkeypatch
):
    monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False)
    monkeypatch.delenv("MLFLOW_EXPERIMENT_ID", raising=False)
    monkeypatch.delenv("MLFLOW_EXPERIMENT_NAME", raising=False)
    monkeypatch.setattr("mlflow.get_tracking_uri", lambda: "file:///tmp/mlruns")

    setup_environment_config(temp_settings_path)

    config = json.loads(temp_settings_path.read_text())
    env_vars = config["env"]
    assert env_vars["MLFLOW_TRACKING_URI"] == "file:///tmp/mlruns"
    assert env_vars["MLFLOW_EXPERIMENT_ID"] == "0"


def test_setup_environment_config_resolves_experiment_name_to_id(temp_settings_path, monkeypatch):
    class DummyExperiment:
        experiment_id = "456"

    class DummyClient:
        def __init__(self, tracking_uri):
            assert tracking_uri == "http://localhost:5000"

        def get_experiment_by_name(self, name):
            assert name == "my-exp"
            return DummyExperiment()

        def create_experiment(self, name):
            raise AssertionError("create_experiment should not be called when experiment exists")

    monkeypatch.setattr("mlflow.tracking.client.MlflowClient", DummyClient)

    setup_environment_config(
        temp_settings_path,
        tracking_uri="http://localhost:5000",
        experiment_name="my-exp",
    )

    config = json.loads(temp_settings_path.read_text())
    env_vars = config["env"]
    assert env_vars["MLFLOW_TRACKING_URI"] == "http://localhost:5000"
    assert env_vars["MLFLOW_EXPERIMENT_ID"] == "456"
    assert env_vars["MLFLOW_EXPERIMENT_NAME"] == "my-exp"