mlflow--mlflow/tests/tracking/test_client.py

import json
import os
import pickle
import threading
import time
import uuid
from pathlib import Path
from unittest import mock
from unittest.mock import Mock, patch

import pytest
from opentelemetry import trace as trace_api
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
from pydantic import BaseModel

import mlflow
from mlflow import MlflowClient, flush_async_logging
from mlflow.config import enable_async_logging
from mlflow.entities import (
    EvaluationDataset,
    ExperimentTag,
    IssueSeverity,
    IssueStatus,
    LoggedModel,
    Run,
    RunInfo,
    RunStatus,
    RunTag,
    SourceType,
    Span,
    SpanLogLevel,
    SpanStatusCode,
    SpanType,
    Trace,
    TraceInfo,
    ViewType,
)
from mlflow.entities.file_info import FileInfo
from mlflow.entities.logged_model_status import LoggedModelStatus
from mlflow.entities.metric import Metric
from mlflow.entities.model_registry import ModelVersion, ModelVersionTag
from mlflow.entities.model_registry.model_version_status import ModelVersionStatus
from mlflow.entities.model_registry.prompt_version import IS_PROMPT_TAG_KEY
from mlflow.entities.param import Param
from mlflow.entities.span import create_mlflow_span
from mlflow.entities.trace_data import TraceData
from mlflow.entities.trace_location import TraceLocation, TraceLocationType, UCSchemaLocation
from mlflow.entities.trace_state import TraceState
from mlflow.entities.trace_status import TraceStatus
from mlflow.environment_variables import MLFLOW_TRACKING_USERNAME
from mlflow.exceptions import (
    MlflowException,
    MlflowNotImplementedException,
    MlflowTraceDataCorrupted,
    MlflowTraceDataNotFound,
)
from mlflow.prompt.registry_utils import PromptCache
from mlflow.store.artifact.artifact_repo import ArtifactRepository
from mlflow.store.entities.paged_list import PagedList
from mlflow.store.model_registry.sqlalchemy_store import (
    SqlAlchemyStore as SqlAlchemyModelRegistryStore,
)
from mlflow.store.tracking import SEARCH_EVALUATION_DATASETS_MAX_RESULTS, SEARCH_MAX_RESULTS_DEFAULT
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore as SqlAlchemyTrackingStore
from mlflow.tracing.constant import SpansLocation, TraceMetadataKey, TraceTagKey
from mlflow.tracing.provider import _get_tracer, trace_disabled
from mlflow.tracing.trace_manager import InMemoryTraceManager
from mlflow.tracing.utils import TraceJSONEncoder
from mlflow.tracking import set_registry_uri
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
from mlflow.tracking._model_registry.utils import (
    _get_store_registry as _get_model_registry_store_registry,
)
from mlflow.tracking._tracking_service.utils import _register, _use_tracking_uri
from mlflow.tracking.default_experiment import DEFAULT_EXPERIMENT_ID
from mlflow.utils.databricks_utils import _construct_databricks_run_url
from mlflow.utils.mlflow_tags import (
    MLFLOW_GIT_COMMIT,
    MLFLOW_PARENT_RUN_ID,
    MLFLOW_PROJECT_ENTRY_POINT,
    MLFLOW_SOURCE_NAME,
    MLFLOW_SOURCE_TYPE,
    MLFLOW_USER,
)
from mlflow.utils.os import is_windows

from tests.tracing.conftest import async_logging_enabled  # noqa: F401
from tests.tracing.helper import create_test_trace_info, get_traces


@pytest.fixture(autouse=True)
def reset_registry_uri():
    yield
    set_registry_uri(None)


@pytest.fixture
def mock_store():
    with mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_get_store:
        mock_store = mock_get_store.return_value
        with mock.patch("mlflow.tracing.client._get_store", return_value=mock_store):
            yield mock_store


@pytest.fixture
def mock_artifact_repo():
    with mock.patch(
        "mlflow.tracking._tracking_service.client.get_artifact_repository"
    ) as mock_get_repo:
        mock_repo = mock_get_repo.return_value
        with mock.patch("mlflow.tracing.client.get_artifact_repository", return_value=mock_repo):
            yield mock_repo


@pytest.fixture
def mock_registry_store():
    mock_store = mock.MagicMock()
    mock_store.create_model_version = mock.create_autospec(
        SqlAlchemyModelRegistryStore.create_model_version
    )
    with mock.patch("mlflow.tracking._model_registry.utils._get_store", return_value=mock_store):
        yield mock_store


@pytest.fixture
def mock_databricks_tracking_store():
    experiment_id = "test-exp-id"
    run_id = "runid"

    class MockDatabricksTrackingStore:
        def __init__(self, run_id, experiment_id):
            self.run_id = run_id
            self.experiment_id = experiment_id

        def get_run(self, *args, **kwargs):
            return Run(
                RunInfo(self.run_id, self.experiment_id, "userid", "status", 0, 1, None), None
            )

    mock_tracking_store = MockDatabricksTrackingStore(run_id, experiment_id)

    with mock.patch(
        "mlflow.tracking._tracking_service.utils._tracking_store_registry.get_store",
        return_value=mock_tracking_store,
    ):
        yield mock_tracking_store


@pytest.fixture
def mock_store_start_trace():
    def _mock_start_trace(trace_info):
        return create_test_trace_info(
            trace_id="tr-123",
            experiment_id=trace_info.experiment_id,
            request_time=trace_info.request_time,
            execution_duration=trace_info.execution_duration,
            state=trace_info.state,
            trace_metadata=trace_info.trace_metadata,
            tags={
                "mlflow.user": "bob",
                "mlflow.artifactLocation": "test",
                **trace_info.tags,
            },
        )

    with mock.patch(
        "mlflow.tracing.client.TracingClient.start_trace", side_effect=_mock_start_trace
    ) as mock_start_trace:
        yield mock_start_trace


@pytest.fixture
def mock_spark_session():
    with mock.patch(
        "mlflow.utils.databricks_utils._get_active_spark_session"
    ) as mock_spark_session:
        yield mock_spark_session.return_value


@pytest.fixture
def mock_time():
    time = 1552319350.244724
    with mock.patch("time.time", return_value=time):
        yield time


@pytest.fixture
def setup_async_logging():
    enable_async_logging(True)
    yield
    flush_async_logging()
    enable_async_logging(False)


def test_client_create_run(mock_store, mock_time):
    experiment_id = mock.Mock()

    MlflowClient().create_run(experiment_id)

    mock_store.create_run.assert_called_once_with(
        experiment_id=experiment_id,
        user_id="unknown",
        start_time=int(mock_time * 1000),
        tags=[],
        run_name=None,
    )


def test_client_create_run_with_name(mock_store, mock_time):
    experiment_id = mock.Mock()

    MlflowClient().create_run(experiment_id, run_name="my name")

    mock_store.create_run.assert_called_once_with(
        experiment_id=experiment_id,
        user_id="unknown",
        start_time=int(mock_time * 1000),
        tags=[],
        run_name="my name",
    )


def test_client_get_trace(mock_store, mock_artifact_repo):
    trace_id = "trace:/catalog.schema/123"
    mock_store.batch_get_traces.return_value = [
        Trace(
            TraceInfo(
                trace_id=trace_id,
                trace_location=TraceLocation(
                    type=TraceLocationType.UC_SCHEMA,
                    uc_schema=UCSchemaLocation(catalog_name="catalog", schema_name="schema"),
                ),
                request_time=123,
                execution_duration=456,
                state=TraceState.OK,
                tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts"},
            ),
            TraceData(
                spans=[
                    Span.from_dict({
                        "name": "predict",
                        "context": {
                            "trace_id": "0x123456789",
                            "span_id": "0x12345",
                        },
                        "parent_id": None,
                        "start_time": 123000000,
                        "end_time": 579000000,
                        "status_code": "OK",
                        "status_message": "",
                        "attributes": {
                            "mlflow.traceRequestId": f'"{trace_id}"',
                            "mlflow.spanType": '"LLM"',
                            "mlflow.spanFunctionName": '"predict"',
                            "mlflow.spanInputs": '{"prompt": "What is the meaning of life?"}',
                            "mlflow.spanOutputs": '{"answer": 42}',
                        },
                        "events": [],
                    })
                ]
            ),
        )
    ]
    trace = MlflowClient().get_trace(trace_id)
    mock_store.batch_get_traces.assert_called_once_with([trace_id], "catalog.schema")
    mock_artifact_repo.download_trace_data.assert_not_called()

    assert trace.info.trace_id == trace_id
    assert trace.info.trace_location.uc_schema.catalog_name == "catalog"
    assert trace.info.trace_location.uc_schema.schema_name == "schema"
    assert trace.info.timestamp_ms == 123
    assert trace.info.execution_time_ms == 456
    assert trace.info.status == TraceStatus.OK
    assert trace.info.tags == {"mlflow.artifactLocation": "dbfs:/path/to/artifacts"}
    assert trace.data.request == '{"prompt": "What is the meaning of life?"}'
    assert trace.data.response == '{"answer": 42}'
    assert len(trace.data.spans) == 1
    assert trace.data.spans[0].name == "predict"
    assert trace.data.spans[0].trace_id == trace_id
    assert trace.data.spans[0].inputs == {"prompt": "What is the meaning of life?"}
    assert trace.data.spans[0].outputs == {"answer": 42}
    assert trace.data.spans[0].start_time_ns == 123000000
    assert trace.data.spans[0].end_time_ns == 579000000
    assert trace.data.spans[0].status.status_code == SpanStatusCode.OK


def test_client_get_trace_empty_result(mock_store):
    mock_store.batch_get_traces.return_value = []
    with pytest.raises(MlflowException, match="not found"):
        MlflowClient().get_trace("trace:/catalog.schema/123")


def test_client_get_trace_from_artifact_repo(mock_store, mock_artifact_repo):
    mock_store.get_trace_info.return_value = TraceInfo(
        trace_id="tr-1234567",
        trace_location=TraceLocation.from_experiment_id("0"),
        request_time=123,
        execution_duration=456,
        state=TraceState.OK,
        tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts"},
    )
    mock_artifact_repo.download_trace_data.return_value = {
        "request": '{"prompt": "What is the meaning of life?"}',
        "response": '{"answer": 42}',
        "spans": [
            {
                "name": "predict",
                "context": {
                    "trace_id": "0x123456789",
                    "span_id": "0x12345",
                },
                "parent_id": None,
                "start_time": 123000000,
                "end_time": 579000000,
                "status_code": "OK",
                "status_message": "",
                "attributes": {
                    "mlflow.traceRequestId": '"tr-1234567"',
                    "mlflow.spanType": '"LLM"',
                    "mlflow.spanFunctionName": '"predict"',
                    "mlflow.spanInputs": '{"prompt": "What is the meaning of life?"}',
                    "mlflow.spanOutputs": '{"answer": 42}',
                },
                "events": [],
            }
        ],
    }
    trace = MlflowClient().get_trace("1234567")
    mock_store.get_trace_info.assert_called_once_with("1234567")
    mock_artifact_repo.download_trace_data.assert_called_once()

    assert trace.info.trace_id == "tr-1234567"
    assert trace.info.experiment_id == "0"
    assert trace.info.timestamp_ms == 123
    assert trace.info.execution_time_ms == 456
    assert trace.info.status == TraceStatus.OK
    assert trace.info.tags == {"mlflow.artifactLocation": "dbfs:/path/to/artifacts"}
    assert trace.data.request == '{"prompt": "What is the meaning of life?"}'
    assert trace.data.response == '{"answer": 42}'
    assert len(trace.data.spans) == 1
    assert trace.data.spans[0].name == "predict"
    assert trace.data.spans[0].trace_id == "tr-1234567"
    assert trace.data.spans[0].inputs == {"prompt": "What is the meaning of life?"}
    assert trace.data.spans[0].outputs == {"answer": 42}
    assert trace.data.spans[0].start_time_ns == 123000000
    assert trace.data.spans[0].end_time_ns == 579000000
    assert trace.data.spans[0].status.status_code == SpanStatusCode.OK


def test_client_get_trace_from_archive_repo(mock_store, mock_artifact_repo):
    trace_info = TraceInfo(
        trace_id="tr-1234567",
        trace_location=TraceLocation.from_experiment_id("0"),
        request_time=123,
        execution_duration=456,
        state=TraceState.OK,
        tags={
            "mlflow.artifactLocation": "dbfs:/path/to/artifacts",
            TraceTagKey.SPANS_LOCATION: SpansLocation.ARCHIVE_REPO,
            TraceTagKey.ARCHIVE_LOCATION: "dbfs:/path/to/archive",
        },
    )
    trace_data = TraceData.from_dict({
        "request": '{"prompt": "What is the meaning of life?"}',
        "response": '{"answer": 42}',
        "spans": [
            {
                "name": "predict",
                "context": {
                    "trace_id": "0x123456789",
                    "span_id": "0x12345",
                },
                "parent_id": None,
                "start_time": 123000000,
                "end_time": 579000000,
                "status_code": "OK",
                "status_message": "",
                "attributes": {
                    "mlflow.traceRequestId": '"tr-1234567"',
                    "mlflow.spanType": '"LLM"',
                    "mlflow.spanFunctionName": '"predict"',
                    "mlflow.spanInputs": '{"prompt": "What is the meaning of life?"}',
                    "mlflow.spanOutputs": '{"answer": 42}',
                },
                "events": [],
            }
        ],
    })
    mock_store.get_trace_info.return_value = trace_info
    mock_store.get_trace.return_value = Trace(info=trace_info, data=trace_data)

    trace = MlflowClient().get_trace("1234567")

    mock_store.get_trace_info.assert_called_once_with("1234567")
    mock_store.get_trace.assert_called_once_with("1234567")
    mock_store.batch_get_traces.assert_not_called()
    mock_artifact_repo.download_archived_trace_data.assert_not_called()
    mock_artifact_repo.download_trace_data.assert_not_called()
    assert trace.info.tags[TraceTagKey.ARCHIVE_LOCATION] == "dbfs:/path/to/archive"
    assert trace.data.spans[0].name == "predict"


def test_client_get_trace_from_archive_repo_returns_empty_spans_when_payload_missing(
    mock_store, mock_artifact_repo
):
    trace_info = TraceInfo(
        trace_id="tr-1234567",
        trace_location=TraceLocation.from_experiment_id("0"),
        request_time=123,
        execution_duration=456,
        state=TraceState.OK,
        tags={
            "mlflow.artifactLocation": "dbfs:/path/to/artifacts",
            TraceTagKey.SPANS_LOCATION: SpansLocation.ARCHIVE_REPO,
            TraceTagKey.ARCHIVE_LOCATION: "dbfs:/path/to/archive",
        },
    )
    mock_store.get_trace_info.return_value = trace_info
    mock_store.get_trace.return_value = Trace(info=trace_info, data=TraceData(spans=[]))

    trace = MlflowClient().get_trace("1234567")

    assert trace.info.trace_id == "tr-1234567"
    assert trace.data.spans == []
    mock_store.get_trace.assert_called_once_with("1234567")
    mock_artifact_repo.download_archived_trace_data.assert_not_called()


def test_client_get_trace_throws_for_missing_or_corrupted_data(mock_store, mock_artifact_repo):
    mock_store.get_trace_info.return_value = TraceInfo(
        trace_id="1234567",
        trace_location=TraceLocation.from_experiment_id("0"),
        request_time=123,
        execution_duration=456,
        state=TraceState.OK,
        tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts"},
    )
    mock_artifact_repo.download_trace_data.side_effect = MlflowTraceDataNotFound("1234567")

    with pytest.raises(
        MlflowException,
        match="Trace with ID 1234567 cannot be loaded because it is missing span data",
    ):
        MlflowClient().get_trace("1234567")

    mock_artifact_repo.download_trace_data.side_effect = MlflowTraceDataCorrupted("1234567")
    with pytest.raises(
        MlflowException,
        match="Trace with ID 1234567 cannot be loaded because its span data is corrupted",
    ):
        MlflowClient().get_trace("1234567")


def test_client_get_trace_throws_for_missing_location_metadata(mock_store, mock_artifact_repo):
    mock_store.get_trace_info.return_value = TraceInfo(
        trace_id="1234567",
        trace_location=TraceLocation.from_experiment_id("0"),
        request_time=123,
        execution_duration=456,
        state=TraceState.OK,
        tags={},
    )

    with pytest.raises(
        MlflowException,
        match="Trace with ID 1234567 cannot be loaded because its span data is corrupted",
    ):
        MlflowClient().get_trace("1234567")

    mock_artifact_repo.download_trace_data.assert_not_called()
    mock_artifact_repo.download_archived_trace_data.assert_not_called()


def test_client_get_trace_from_archive_repo_does_not_require_archive_location_tag(
    mock_store, mock_artifact_repo
):
    trace_info = TraceInfo(
        trace_id="tr-1234567",
        trace_location=TraceLocation.from_experiment_id("0"),
        request_time=123,
        execution_duration=456,
        state=TraceState.OK,
        tags={TraceTagKey.SPANS_LOCATION: SpansLocation.ARCHIVE_REPO},
    )
    mock_store.get_trace_info.return_value = trace_info
    mock_store.get_trace.return_value = Trace(info=trace_info, data=TraceData(spans=[]))

    trace = MlflowClient().get_trace("1234567")

    assert trace.info.trace_id == "tr-1234567"
    assert trace.data.spans == []
    mock_store.get_trace.assert_called_once_with("1234567")
    mock_artifact_repo.download_trace_data.assert_not_called()
    mock_artifact_repo.download_archived_trace_data.assert_not_called()


@pytest.mark.parametrize("include_spans", [True, False])
@pytest.mark.parametrize("num_results", [0, 5])
def test_client_search_traces_with_get_traces(
    mock_store, mock_artifact_repo, include_spans, num_results
):
    mock_trace_infos = [
        TraceInfo(
            trace_id=f"trace:/catalog.schema/{i}",
            trace_location=TraceLocation(
                type=TraceLocationType.UC_SCHEMA,
                uc_schema=UCSchemaLocation(catalog_name="catalog", schema_name="schema"),
            ),
            request_time=123,
            execution_duration=456,
            state=TraceState.OK,
        )
        for i in range(num_results)
    ]
    mock_store.search_traces.return_value = (mock_trace_infos, None)
    mock_store.batch_get_traces.return_value = [
        Trace(info=info, data=TraceData(spans=[])) for info in mock_trace_infos
    ]

    results = MlflowClient().search_traces(
        locations=["catalog.schema"],
        include_spans=include_spans,
    )
    mock_store.search_traces.assert_called_once_with(
        experiment_ids=None,
        filter_string=None,
        max_results=100,
        order_by=None,
        page_token=None,
        model_id=None,
        locations=["catalog.schema"],
    )
    assert len(results) == num_results

    if include_spans and num_results > 0:
        mock_store.batch_get_traces.assert_called_once_with(
            [f"trace:/catalog.schema/{i}" for i in range(num_results)],
            "catalog.schema",
        )
    else:
        mock_store.batch_get_traces.assert_not_called()

    mock_artifact_repo.download_trace_data.assert_not_called()

    # The TraceInfo is already fetched prior to the upload_trace_data call,
    # so we should not call _get_trace_info again
    mock_store.get_trace_info.assert_not_called()


def test_client_search_traces_with_large_results(mock_store, mock_artifact_repo):
    mock_trace_infos = [
        TraceInfo(
            trace_id=f"trace:/catalog.schema/{i}",
            trace_location=TraceLocation(
                type=TraceLocationType.UC_SCHEMA,
                uc_schema=UCSchemaLocation(catalog_name="catalog", schema_name="schema"),
            ),
            request_time=123,
            execution_duration=456,
            state=TraceState.OK,
        )
        for i in range(100)
    ]
    mock_store.search_traces.return_value = (mock_trace_infos, None)

    mock_store.batch_get_traces.return_value = [
        Trace(info=mock_trace_infos[0], data=TraceData(spans=[])) for i in range(10)
    ]

    results = MlflowClient().search_traces(locations=["catalog.schema"])
    mock_store.search_traces.assert_called_once_with(
        experiment_ids=None,
        filter_string=None,
        max_results=100,
        order_by=None,
        page_token=None,
        model_id=None,
        locations=["catalog.schema"],
    )
    assert len(results) == 100
    assert mock_store.batch_get_traces.call_count == 10
    assert mock_store.batch_get_traces.has_calls([
        mock.call([f"trace:/catalog.schema/{j * 10 + i}" for i in range(10)], "catalog.schema")
        for j in range(10)
    ])
    mock_artifact_repo.download_trace_data.assert_not_called()


@pytest.mark.parametrize("include_spans", [True, False])
def test_client_search_traces_mixed(mock_store, mock_artifact_repo, include_spans):
    mock_traces = [
        TraceInfo(
            trace_id="1234567",
            trace_location=TraceLocation(
                type=TraceLocationType.UC_SCHEMA,
                uc_schema=UCSchemaLocation(catalog_name="catalog", schema_name="schema"),
            ),
            request_time=123,
            execution_duration=456,
            state=TraceState.OK,
            tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts/1"},
        ),
        TraceInfo(
            trace_id="8910",
            trace_location=TraceLocation.from_experiment_id("1"),
            request_time=456,
            execution_duration=789,
            state=TraceState.OK,
            tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts/2"},
        ),
    ]
    mock_store.search_traces.return_value = (mock_traces, None)
    mock_store.batch_get_traces.return_value = [
        Trace(info=mock_traces[0], data=TraceData(spans=[]))
    ]
    mock_artifact_repo.download_trace_data.return_value = {}
    results = MlflowClient().search_traces(
        locations=["1", "catalog.schema"], include_spans=include_spans
    )

    mock_store.search_traces.assert_called_once_with(
        experiment_ids=None,
        filter_string=None,
        max_results=100,
        order_by=None,
        page_token=None,
        model_id=None,
        locations=["1", "catalog.schema"],
    )
    assert len(results) == 2
    if include_spans:
        mock_store.batch_get_traces.assert_called_once_with(["1234567"], "catalog.schema")
        mock_artifact_repo.download_trace_data.assert_called()
    else:
        mock_store.batch_get_traces.assert_not_called()
        mock_artifact_repo.download_trace_data.assert_not_called()


@pytest.mark.parametrize("include_spans", [True, False])
@pytest.mark.parametrize("num_results", [0, 5])
def test_client_search_traces_with_get_traces_tracking_store(
    mock_store, mock_artifact_repo, include_spans, num_results
):
    mock_trace_infos = [
        TraceInfo(
            trace_id=f"tr-123456789{i}",
            trace_location=TraceLocation.from_experiment_id(f"exp-{i}"),
            request_time=123,
            execution_duration=456,
            state=TraceState.OK,
            tags={TraceTagKey.SPANS_LOCATION: SpansLocation.TRACKING_STORE},
        )
        for i in range(num_results)
    ]
    mock_store.search_traces.return_value = (mock_trace_infos, None)
    mock_store.batch_get_traces.return_value = [
        Trace(info=info, data=TraceData(spans=[])) for info in mock_trace_infos
    ]

    results = MlflowClient().search_traces(
        locations=["exp-0", "exp-1", "exp-2"],
        include_spans=include_spans,
    )
    mock_store.search_traces.assert_called_once_with(
        experiment_ids=None,
        filter_string=None,
        max_results=100,
        order_by=None,
        page_token=None,
        model_id=None,
        locations=["exp-0", "exp-1", "exp-2"],
    )
    assert len(results) == num_results

    if include_spans and num_results > 0:
        mock_store.batch_get_traces.assert_called_once_with(
            [f"tr-123456789{i}" for i in range(num_results)],
            None,
        )
    else:
        mock_store.batch_get_traces.assert_not_called()

    mock_artifact_repo.download_trace_data.assert_not_called()

    # The TraceInfo is already fetched prior to the upload_trace_data call,
    # so we should not call _get_trace_info again
    mock_store.get_trace_info.assert_not_called()


@pytest.mark.parametrize("include_spans", [True, False])
def test_client_search_traces_with_artifact_repo(mock_store, mock_artifact_repo, include_spans):
    mock_traces = [
        TraceInfo(
            trace_id="tr-1234567",
            trace_location=TraceLocation.from_experiment_id("1"),
            request_time=123,
            execution_duration=456,
            state=TraceState.OK,
            tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts/1"},
        ),
        TraceInfo(
            trace_id="tr-8910",
            trace_location=TraceLocation.from_experiment_id("2"),
            request_time=456,
            execution_duration=789,
            state=TraceState.OK,
            tags={"mlflow.artifactLocation": "dbfs:/path/to/artifacts/2"},
        ),
    ]
    mock_store.search_traces.return_value = (mock_traces, None)
    mock_artifact_repo.download_trace_data.return_value = {}
    results = MlflowClient().search_traces(locations=["1", "2", "3"], include_spans=include_spans)

    mock_store.search_traces.assert_called_once_with(
        experiment_ids=None,
        filter_string=None,
        max_results=100,
        order_by=None,
        page_token=None,
        model_id=None,
        locations=["1", "2", "3"],
    )
    assert len(results) == 2
    if include_spans:
        mock_artifact_repo.download_trace_data.assert_called()
    else:
        mock_artifact_repo.download_trace_data.assert_not_called()

    # The TraceInfo is already fetched prior to the upload_trace_data call,
    # so we should not call _get_trace_info again
    mock_store.get_trace_info.assert_not_called()


@pytest.mark.parametrize("include_spans", [True, False])
def test_client_search_traces_trace_data_download_error(mock_store, include_spans):
    class CustomArtifactRepository(ArtifactRepository):
        def log_artifact(self, local_file, artifact_path=None):
            raise NotImplementedError("Should not be called")

        def log_artifacts(self, local_dir, artifact_path=None):
            raise NotImplementedError("Should not be called")

        def list_artifacts(self, path):
            raise NotImplementedError("Should not be called")

        def _download_file(self, *args, **kwargs):
            raise Exception("Failed to download trace data")

    with mock.patch(
        "mlflow.tracing.client.get_artifact_repository",
        return_value=CustomArtifactRepository("test"),
    ) as mock_get_artifact_repository:
        mock_traces = [
            TraceInfo(
                trace_id="1234567",
                trace_location=TraceLocation.from_experiment_id("1"),
                request_time=123,
                execution_duration=456,
                state=TraceState.OK,
                tags={"mlflow.artifactLocation": "test"},
            ),
        ]
        mock_store.search_traces.return_value = (mock_traces, None)
        traces = MlflowClient().search_traces(locations=["1"], include_spans=include_spans)

        if include_spans:
            assert traces == []
            mock_get_artifact_repository.assert_called()
        else:
            assert len(traces) == 1
            assert traces[0].info.trace_id == "1234567"
            mock_get_artifact_repository.assert_not_called()


def test_client_search_traces_validates_experiment_ids_type():
    with pytest.raises(MlflowException, match=r"locations must be a list"):
        MlflowClient().search_traces(locations=4)

    with pytest.raises(MlflowException, match=r"locations must be a list"):
        MlflowClient().search_traces(locations="4")


def test_client_delete_traces(mock_store):
    MlflowClient().delete_traces(
        experiment_id="0",
        max_timestamp_millis=1,
        max_traces=2,
        trace_ids=["tr-1234"],
    )
    mock_store.delete_traces.assert_called_once_with(
        experiment_id="0",
        max_timestamp_millis=1,
        max_traces=2,
        trace_ids=["tr-1234"],
    )


@pytest.fixture
def disable_prompt_cache():
    from mlflow.environment_variables import (
        MLFLOW_ALIAS_PROMPT_CACHE_TTL_SECONDS,
        MLFLOW_VERSION_PROMPT_CACHE_TTL_SECONDS,
    )

    MLFLOW_ALIAS_PROMPT_CACHE_TTL_SECONDS.set(0)
    MLFLOW_VERSION_PROMPT_CACHE_TTL_SECONDS.set(0)
    yield
    MLFLOW_ALIAS_PROMPT_CACHE_TTL_SECONDS.unset()
    MLFLOW_VERSION_PROMPT_CACHE_TTL_SECONDS.unset()


@pytest.fixture(autouse=True)
def reset_prompt_cache():
    PromptCache._reset_instance()
    yield
    PromptCache._reset_instance()


@pytest.fixture(params=["file", "sqlalchemy"])
def tracking_uri(request, tmp_path, db_uri):
    """Set an MLflow Tracking URI with different type of backend."""
    if request.param == "file":
        pytest.skip("FileStore is no longer supported.")
    if "MLFLOW_SKINNY" in os.environ and request.param == "sqlalchemy":
        pytest.skip("SQLAlchemy store is not available in skinny.")

    original_tracking_uri = mlflow.get_tracking_uri()

    if request.param == "file":
        tracking_uri = tmp_path.joinpath("file").as_uri()
    elif request.param == "sqlalchemy":
        tracking_uri = db_uri

    # NB: MLflow tracer does not handle the change of tracking URI well,
    # so we need to reset the tracer to switch the tracking URI during testing.
    mlflow.tracing.disable()
    mlflow.set_tracking_uri(tracking_uri)
    mlflow.tracing.enable()

    yield tracking_uri

    # Reset tracking URI
    mlflow.set_tracking_uri(original_tracking_uri)


@pytest.mark.parametrize("with_active_run", [True, False])
def test_start_and_end_trace(tracking_uri, with_active_run, async_logging_enabled):
    client = MlflowClient(tracking_uri)

    experiment_id = client.create_experiment("test_experiment")

    class TestModel:
        def predict(self, x, y):
            root_span = client.start_trace(
                name="predict",
                inputs={"x": x, "y": y},
                tags={"tag": "tag_value"},
                experiment_id=experiment_id,
            )
            trace_id = root_span.trace_id

            z = x + y

            child_span = client.start_span(
                "child_span_1",
                span_type=SpanType.LLM,
                trace_id=trace_id,
                parent_id=root_span.span_id,
                inputs={"z": z},
            )

            z = z + 2

            client.end_span(
                trace_id=trace_id,
                span_id=child_span.span_id,
                outputs={"output": z},
                attributes={"delta": 2},
            )

            res = self.square(z, trace_id, root_span.span_id)
            client.end_trace(trace_id, outputs={"output": res}, status="OK")
            return res

        def square(self, t, trace_id, parent_id):
            span = client.start_span(
                "child_span_2",
                trace_id=trace_id,
                parent_id=parent_id,
                inputs={"t": t},
            )

            res = t**2
            time.sleep(0.1)

            client.end_span(
                trace_id=trace_id,
                span_id=span.span_id,
                outputs={"output": res},
            )
            return res

    model = TestModel()
    if with_active_run:
        with mlflow.start_run(experiment_id=experiment_id) as run:
            model.predict(1, 2)
            run_id = run.info.run_id
    else:
        model.predict(1, 2)

    if async_logging_enabled:
        mlflow.flush_trace_async_logging(terminate=True)

    trace_id = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True).info.trace_id

    # Validate that trace is logged to the backend
    trace = client.get_trace(trace_id)
    assert trace is not None

    assert trace.info.trace_id is not None
    assert trace.info.execution_time_ms >= 0.1 * 1e3  # at least 0.1 sec
    assert trace.info.status == TraceStatus.OK
    assert trace.info.request_metadata[TraceMetadataKey.INPUTS] == '{"x": 1, "y": 2}'
    assert trace.info.request_metadata[TraceMetadataKey.OUTPUTS] == '{"output": 25}'
    if with_active_run:
        assert trace.info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_id
        assert trace.info.experiment_id == run.info.experiment_id
    else:
        assert trace.info.experiment_id == experiment_id

    assert trace.data.request == '{"x": 1, "y": 2}'
    assert trace.data.response == '{"output": 25}'
    assert len(trace.data.spans) == 3

    span_name_to_span = {span.name: span for span in trace.data.spans}
    root_span = span_name_to_span["predict"]
    # NB: Start time of root span and trace info does not match because there is some
    #   latency for starting the trace within the backend
    # assert root_span.start_time_ns // 1e6 == trace.info.timestamp_ms
    assert root_span.parent_id is None
    assert root_span.attributes == {
        "mlflow.experimentId": experiment_id,
        "mlflow.traceRequestId": trace.info.trace_id,
        "mlflow.spanType": "UNKNOWN",
        "mlflow.spanLogLevel": SpanLogLevel.DEBUG,
        "mlflow.spanInputs": {"x": 1, "y": 2},
        "mlflow.spanOutputs": {"output": 25},
    }

    child_span_1 = span_name_to_span["child_span_1"]
    assert child_span_1.parent_id == root_span.span_id
    assert child_span_1.attributes == {
        "mlflow.traceRequestId": trace.info.trace_id,
        "mlflow.spanType": "LLM",
        "mlflow.spanLogLevel": SpanLogLevel.INFO,
        "mlflow.spanInputs": {"z": 3},
        "mlflow.spanOutputs": {"output": 5},
        "delta": 2,
    }

    child_span_2 = span_name_to_span["child_span_2"]
    assert child_span_2.parent_id == root_span.span_id
    assert child_span_2.attributes == {
        "mlflow.traceRequestId": trace.info.trace_id,
        "mlflow.spanType": "UNKNOWN",
        "mlflow.spanLogLevel": SpanLogLevel.DEBUG,
        "mlflow.spanInputs": {"t": 5},
        "mlflow.spanOutputs": {"output": 25},
    }
    assert child_span_2.start_time_ns <= child_span_2.end_time_ns - 0.1 * 1e6


def test_start_trace_with_run_id(tracking_uri, async_logging_enabled):
    client = MlflowClient(tracking_uri)

    experiment_id = client.create_experiment(f"test_experiment_{uuid.uuid4().hex}")
    run = client.create_run(experiment_id=experiment_id)

    root_span = client.start_trace(
        name="test",
        experiment_id=experiment_id,
        run_id=run.info.run_id,
    )
    client.end_trace(root_span.trace_id)

    traces = client.search_traces(
        locations=[experiment_id],
        include_spans=False,
        flush=True,
    )

    assert len(traces) == 1
    trace_info = traces[0].info
    assert trace_info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run.info.run_id
    assert trace_info.experiment_id == experiment_id


def test_start_and_end_trace_capture_falsy_input_and_output(tracking_uri):
    # This test is to verify that falsy input and output values are correctly logged
    client = MlflowClient(tracking_uri)
    experiment_id = client.create_experiment("test_experiment")

    root = client.start_trace(name="root", experiment_id=experiment_id, inputs=[])
    span = client.start_span(name="child", trace_id=root.trace_id, parent_id=root.span_id, inputs=0)
    client.end_span(trace_id=root.trace_id, span_id=span.span_id, outputs=False)
    client.end_trace(trace_id=root.trace_id, outputs="")

    trace = client.get_trace(root.trace_id, flush=True)
    assert trace.data.spans[0].inputs == []
    assert trace.data.spans[0].outputs == ""
    assert trace.data.spans[1].inputs == 0
    assert trace.data.spans[1].outputs is False


# TODO: we should investigate whether we need to support this
@pytest.mark.skip(reason="This is not supported by latest span-level export")
@pytest.mark.usefixtures("reset_active_experiment")
def test_start_and_end_trace_before_all_span_end(async_logging_enabled):
    # This test is to verify that the trace is still exported even if some spans are not ended
    exp_id = mlflow.set_experiment("test_experiment_1").experiment_id

    class TestModel:
        def __init__(self):
            self._client = MlflowClient()

        def predict(self, x):
            root_span = self._client.start_trace(name="predict")
            trace_id = root_span.trace_id
            child_span = self._client.start_span(
                "ended-span",
                trace_id=trace_id,
                parent_id=root_span.span_id,
            )
            time.sleep(0.1)
            self._client.end_span(trace_id, child_span.span_id)

            res = self.square(x, trace_id, root_span.span_id)
            self._client.end_trace(trace_id)
            return res

        def square(self, t, trace_id, parent_id):
            self._client.start_span("non-ended-span", trace_id=trace_id, parent_id=parent_id)
            time.sleep(0.1)
            # The span created above is not ended
            return t**2

    model = TestModel()
    model.predict(1)

    if async_logging_enabled:
        mlflow.flush_trace_async_logging(terminate=True)

    traces = MlflowClient().search_traces(locations=[exp_id])
    assert len(traces) == 1

    trace_info = traces[0].info
    assert trace_info.trace_id is not None
    assert trace_info.experiment_id == exp_id
    assert trace_info.timestamp_ms is not None
    assert trace_info.execution_time_ms is not None
    assert trace_info.status == TraceStatus.OK

    trace_data = traces[0].data
    assert trace_data.request is None
    assert trace_data.response is None
    assert len(trace_data.spans) == 3  # The non-ended span should be also included in the trace

    span_name_to_span = {span.name: span for span in trace_data.spans}
    root_span = span_name_to_span["predict"]
    assert root_span.parent_id is None
    assert root_span.status.status_code == SpanStatusCode.OK

    ended_span = span_name_to_span["ended-span"]
    assert ended_span.parent_id == root_span.span_id
    assert ended_span.start_time_ns < ended_span.end_time_ns
    assert ended_span.status.status_code == SpanStatusCode.OK

    # The non-ended span should have null end_time and UNSET status
    non_ended_span = span_name_to_span["non-ended-span"]
    assert non_ended_span.parent_id == root_span.span_id
    assert non_ended_span.start_time_ns is not None
    assert non_ended_span.end_time_ns is None
    assert non_ended_span.status.status_code == SpanStatusCode.UNSET


def test_log_trace_with_databricks_tracking_uri(mock_store_start_trace, monkeypatch):
    monkeypatch.setenv("MLFLOW_EXPERIMENT_NAME", "test")
    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "bob")

    mock_experiment = mock.MagicMock()
    mock_experiment.experiment_id = "test_experiment_id"

    class TestModel:
        def __init__(self):
            self._client = MlflowClient()

        def predict(self, x, y):
            root_span = self._client.start_trace(
                name="predict",
                inputs={"x": x, "y": y},
                # Trying to override mlflow.user tag, which will be ignored
                tags={"tag": "tag_value", "mlflow.user": "unknown"},
            )
            trace_id = root_span.trace_id

            z = x + y

            child_span = self._client.start_span(
                "child_span_1",
                span_type=SpanType.LLM,
                trace_id=trace_id,
                parent_id=root_span.span_id,
                inputs={"z": z},
            )

            z = z + 2

            self._client.end_span(
                trace_id=trace_id,
                span_id=child_span.span_id,
                outputs={"output": z},
                attributes={"delta": 2},
            )
            self._client.end_trace(trace_id, outputs=z, status="OK")
            return z

    model = TestModel()

    with (
        mock.patch("mlflow.get_tracking_uri", return_value="databricks"),
        mock.patch("mlflow.tracking.context.default_context._get_source_name", return_value="test"),
        mock.patch(
            "mlflow.tracing.client.TracingClient._upload_trace_data"
        ) as mock_upload_trace_data,
        mock.patch(
            "mlflow.tracing.client.TracingClient.set_trace_tags",
        ),
        mock.patch(
            "mlflow.tracing.client.TracingClient.set_trace_tag",
        ),
        mock.patch(
            "mlflow.tracing.client.TracingClient.get_trace_info",
        ),
        mock.patch(
            "mlflow.MlflowClient.get_experiment_by_name",
            return_value=mock_experiment,
        ),
    ):
        model.predict(1, 2)
        mlflow.flush_trace_async_logging(terminate=True)

    mock_store_start_trace.assert_called_once()
    mock_upload_trace_data.assert_called_once()


def test_start_and_end_trace_does_not_log_trace_when_disabled(
    tracking_uri, monkeypatch, async_logging_enabled
):
    client = MlflowClient(tracking_uri)
    experiment_id = client.create_experiment("test_experiment")

    @trace_disabled
    def func():
        span = client.start_trace(
            name="predict",
            experiment_id=experiment_id,
            inputs={"x": 1, "y": 2},
            attributes={"attr": "value"},
            tags={"tag": "tag_value"},
        )
        child_span = client.start_span(
            "child_span_1",
            trace_id=span.trace_id,
            parent_id=span.span_id,
        )
        client.end_span(
            trace_id=span.trace_id,
            span_id=child_span.span_id,
            outputs={"output": 5},
        )
        client.end_trace(span.trace_id, outputs=5, status="OK")
        return "done"

    mock_logger = mock.MagicMock()
    monkeypatch.setattr(mlflow.tracking.client, "_logger", mock_logger)

    res = func()

    assert res == "done"
    assert client.search_traces(locations=[experiment_id]) == []
    # No warning should be issued
    mock_logger.warning.assert_not_called()


def test_start_trace_within_active_run(async_logging_enabled):
    exp_id = mlflow.create_experiment("test")

    client = mlflow.MlflowClient()
    with mlflow.start_run():
        root_span = client.start_trace(
            name="test",
            experiment_id=exp_id,
        )
        client.end_trace(root_span.trace_id)

    if async_logging_enabled:
        mlflow.flush_trace_async_logging(terminate=True)

    traces = client.search_traces(locations=[exp_id])
    assert len(traces) == 1
    assert traces[0].info.experiment_id == exp_id


def test_start_trace_raise_error_when_active_trace_exists():
    with mlflow.start_span("fluent_span"):
        with pytest.raises(MlflowException, match=r"Another trace is already set in the global"):
            mlflow.tracking.MlflowClient().start_trace("test")


def test_end_trace_raise_error_when_trace_not_exist():
    client = mlflow.tracking.MlflowClient()
    mock_tracing_client = mock.MagicMock()
    mock_tracing_client.get_trace.return_value = None
    client._tracing_client = mock_tracing_client

    with pytest.raises(MlflowException, match=r"Trace with ID test not found"):
        client.end_trace("test")


def test_end_trace_works_for_trace_in_pending_status():
    client = mlflow.tracking.MlflowClient()
    mock_tracing_client = mock.MagicMock()
    mock_tracing_client.get_trace.return_value = Trace(
        info=create_test_trace_info("test", state=TraceState.IN_PROGRESS), data=TraceData()
    )
    client._tracing_client = mock_tracing_client
    client.end_span = lambda *args: None

    client.end_trace("test")


@pytest.mark.parametrize("state", [TraceState.OK, TraceState.ERROR])
def test_end_trace_raise_error_for_trace_in_end_status(state):
    client = mlflow.tracking.MlflowClient()
    mock_tracing_client = mock.MagicMock()
    mock_tracing_client.get_trace.return_value = Trace(
        info=create_test_trace_info("test", state=state), data=TraceData()
    )
    client._tracing_client = mock_tracing_client

    with pytest.raises(MlflowException, match=r"Trace with ID test already finished"):
        client.end_trace("test")


def test_trace_status_either_pending_or_end():
    all_statuses = {status.value for status in TraceStatus}
    pending_or_end_statuses = TraceStatus.pending_statuses() | TraceStatus.end_statuses()
    unclassified_statuses = all_statuses - pending_or_end_statuses
    assert len(unclassified_statuses) == 0, (
        f"Please add {unclassified_statuses} to "
        "either pending_statuses or end_statuses in TraceStatus class definition"
    )


def test_start_span_raise_error_when_parent_id_is_not_provided():
    with pytest.raises(MlflowException, match=r"start_span\(\) must be called with"):
        mlflow.tracking.MlflowClient().start_span("span_name", trace_id="test", parent_id=None)


def test_log_trace(tracking_uri):
    client = MlflowClient(tracking_uri)
    experiment_id = client.create_experiment("test_experiment")

    span = client.start_trace(
        name="test",
        span_type=SpanType.LLM,
        experiment_id=experiment_id,
        tags={"custom_tag": "tag_value"},
    )
    client.end_trace(span.trace_id, status="OK")

    trace = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True)

    # Purge all traces in the backend once
    client.delete_traces(experiment_id=experiment_id, trace_ids=[trace.info.trace_id])
    assert client.search_traces(locations=[experiment_id]) == []

    # Log the trace manually — _log_trace triggers async export via span processor
    new_trace_id = client._log_trace(trace)

    # Validate the trace is added to the backend (flush=True waits for async writes)
    backend_traces = client.search_traces(locations=[experiment_id], flush=True)
    assert len(backend_traces) == 1
    assert backend_traces[0].info.trace_id == new_trace_id  # new request ID is assigned
    assert backend_traces[0].info.experiment_id == experiment_id
    assert backend_traces[0].info.status == trace.info.status
    assert backend_traces[0].info.tags["custom_tag"] == "tag_value"
    assert backend_traces[0].info.request_preview == trace.info.request_preview
    assert backend_traces[0].info.response_preview == trace.info.response_preview
    assert len(backend_traces[0].data.spans) == len(trace.data.spans)
    assert backend_traces[0].data.spans[0].name == trace.data.spans[0].name

    # If the experiment ID is None in the given trace, it should be set to the default experiment
    trace.info.experiment_id = None
    new_trace_id = client._log_trace(trace)
    backend_traces = client.search_traces(locations=[DEFAULT_EXPERIMENT_ID], flush=True)
    assert len(backend_traces) == 1


@pytest.mark.filterwarnings("ignore::FutureWarning")
def test_search_traces_experiment_ids_deprecation_warning():
    client = MlflowClient()
    exp_id = mlflow.set_experiment("test_experiment_deprecation").experiment_id

    # Test that using experiment_ids shows a deprecation warning
    with pytest.warns(FutureWarning, match="experiment_ids.*deprecated.*use.*locations"):
        result = client.search_traces(experiment_ids=[exp_id])
    assert isinstance(result, list)


def test_ignore_exception_from_tracing_logic(monkeypatch, async_logging_enabled):
    exp_id = mlflow.set_experiment("test_experiment_1").experiment_id
    client = MlflowClient()

    class TestModel:
        def predict(self, x):
            root_span = client.start_trace(experiment_id=exp_id, name="predict")
            trace_id = root_span.trace_id
            child_span = client.start_span(
                name="child", trace_id=trace_id, parent_id=root_span.span_id
            )
            client.end_span(trace_id, child_span.span_id)
            client.end_trace(trace_id)
            return x

    model = TestModel()

    # Mock the span processor's on_end handler to raise an exception
    processor = _get_tracer(__name__).span_processor

    def _always_fail(*args, **kwargs):
        raise ValueError("Some error")

    # Exception while starting the trace should be caught not raise
    monkeypatch.setattr(processor, "on_start", _always_fail)
    response = model.predict(1)
    assert response == 1
    assert len(get_traces()) == 0

    # Exception while ending the trace should be caught not raise
    monkeypatch.setattr(processor, "on_end", _always_fail)
    response = model.predict(1)
    assert response == 1
    assert len(get_traces()) == 0


def test_set_and_delete_trace_tag_on_active_trace(monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "bob")
    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")

    client = mlflow.tracking.MlflowClient()

    root_span = client.start_trace(name="test")
    trace_id = root_span.trace_id
    client.set_trace_tag(trace_id, "foo", "bar")
    client.end_trace(trace_id)

    trace = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True)
    assert trace.info.tags["foo"] == "bar"


def test_set_trace_tag_on_logged_trace(mock_store):
    mlflow.tracking.MlflowClient().set_trace_tag("test", "foo", "bar")
    mlflow.tracking.MlflowClient().set_trace_tag("test", "mlflow.some.reserved.tag", "value")
    mock_store.set_trace_tag.assert_has_calls([
        mock.call("test", "foo", "bar"),
        mock.call("test", "mlflow.some.reserved.tag", "value"),
    ])


@pytest.mark.parametrize(
    "key",
    [
        TraceTagKey.SPANS_LOCATION,
        TraceTagKey.ARCHIVE_LOCATION,
        TraceTagKey.ARCHIVAL_FAILURE,
    ],
)
def test_set_trace_tag_skips_immutable_internal_tags_on_active_trace(monkeypatch, key):
    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "bob")
    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")

    client = mlflow.tracking.MlflowClient()
    root_span = client.start_trace(name="test")
    trace_id = root_span.trace_id

    with patch("mlflow.tracing.client._logger") as mock_logger:
        client.set_trace_tag(trace_id, key, "s3://bucket/archive/test")

    client.end_trace(trace_id)

    trace = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True)
    if key == TraceTagKey.SPANS_LOCATION:
        assert trace.info.tags[key] == SpansLocation.TRACKING_STORE.value
    else:
        assert key not in trace.info.tags
    mock_logger.warning.assert_called_once_with(
        f"Tag '{key}' is immutable and cannot be set on a trace."
    )


@pytest.mark.parametrize(
    "key",
    [
        TraceTagKey.SPANS_LOCATION,
        TraceTagKey.ARCHIVE_LOCATION,
        TraceTagKey.ARCHIVAL_FAILURE,
    ],
)
def test_set_trace_tag_skips_immutable_internal_tags(mock_store, key):
    with patch("mlflow.tracing.client._logger") as mock_logger:
        mlflow.tracking.MlflowClient().set_trace_tag("test", key, "s3://bucket/archive/test")

    mock_store.set_trace_tag.assert_not_called()
    mock_logger.warning.assert_called_once_with(
        f"Tag '{key}' is immutable and cannot be set on a trace."
    )


def test_delete_trace_tag_on_active_trace(monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "bob")
    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")

    client = mlflow.tracking.MlflowClient()
    root_span = client.start_trace(name="test", tags={"foo": "bar", "baz": "qux"})
    trace_id = root_span.trace_id
    client.delete_trace_tag(trace_id, "foo")
    client.end_trace(trace_id)

    trace = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True)
    assert "baz" in trace.info.tags
    assert "foo" not in trace.info.tags


def test_delete_trace_tag_on_logged_trace(mock_store):
    mlflow.tracking.MlflowClient().delete_trace_tag("test", "foo")
    mock_store.delete_trace_tag.assert_called_once_with("test", "foo")


@pytest.mark.parametrize(
    "key",
    [
        TraceTagKey.SPANS_LOCATION,
        TraceTagKey.ARCHIVE_LOCATION,
    ],
)
def test_delete_trace_tag_skips_immutable_internal_tags_on_active_trace(monkeypatch, key):
    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "bob")
    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")

    client = mlflow.tracking.MlflowClient()
    root_span = client.start_trace(name="test", tags={"foo": "bar"})
    trace_id = root_span.trace_id

    with patch("mlflow.tracing.client._logger") as mock_logger:
        client.delete_trace_tag(trace_id, key)

    client.end_trace(trace_id)

    trace = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True)
    assert trace.info.tags["foo"] == "bar"
    if key == TraceTagKey.SPANS_LOCATION:
        assert trace.info.tags[key] == SpansLocation.TRACKING_STORE.value
    else:
        assert key not in trace.info.tags
    mock_logger.warning.assert_called_once_with(
        f"Tag '{key}' is immutable and cannot be deleted on a trace."
    )


@pytest.mark.parametrize(
    "key",
    [
        TraceTagKey.SPANS_LOCATION,
        TraceTagKey.ARCHIVE_LOCATION,
    ],
)
def test_delete_trace_tag_skips_immutable_internal_tags(mock_store, key):
    with patch("mlflow.tracing.client._logger") as mock_logger:
        mlflow.tracking.MlflowClient().delete_trace_tag("test", key)

    mock_store.delete_trace_tag.assert_not_called()
    mock_logger.warning.assert_called_once_with(
        f"Tag '{key}' is immutable and cannot be deleted on a trace."
    )


def test_delete_trace_tag_allows_clearing_archival_failure_on_active_trace(monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "bob")
    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")

    client = mlflow.tracking.MlflowClient()
    root_span = client.start_trace(name="test", tags={"foo": "bar"})
    trace_id = root_span.trace_id
    with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
        trace.info.tags[TraceTagKey.ARCHIVAL_FAILURE] = "MALFORMED_TRACE"

    client.delete_trace_tag(trace_id, TraceTagKey.ARCHIVAL_FAILURE)
    client.end_trace(trace_id)

    trace = mlflow.get_trace(mlflow.get_last_active_trace_id(), flush=True)
    assert trace.info.tags["foo"] == "bar"
    assert TraceTagKey.ARCHIVAL_FAILURE not in trace.info.tags


def test_delete_trace_tag_allows_clearing_archival_failure_on_logged_trace(mock_store):
    mlflow.tracking.MlflowClient().delete_trace_tag("test", TraceTagKey.ARCHIVAL_FAILURE)

    mock_store.delete_trace_tag.assert_called_once_with("test", TraceTagKey.ARCHIVAL_FAILURE)


def test_client_create_experiment(mock_store):
    MlflowClient().create_experiment("someName", "someLocation", {"key1": "val1", "key2": "val2"})

    mock_store.create_experiment.assert_called_once_with(
        artifact_location="someLocation",
        tags=[ExperimentTag("key1", "val1"), ExperimentTag("key2", "val2")],
        name="someName",
    )


def test_client_create_run_overrides(mock_store):
    experiment_id = mock.Mock()
    user = mock.Mock()
    start_time = mock.Mock()
    run_name = mock.Mock()
    tags = {
        MLFLOW_USER: user,
        MLFLOW_PARENT_RUN_ID: mock.Mock(),
        MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB),
        MLFLOW_SOURCE_NAME: mock.Mock(),
        MLFLOW_PROJECT_ENTRY_POINT: mock.Mock(),
        MLFLOW_GIT_COMMIT: mock.Mock(),
        "other-key": "other-value",
    }

    MlflowClient().create_run(experiment_id, start_time, tags, run_name)

    mock_store.create_run.assert_called_once_with(
        experiment_id=experiment_id,
        user_id=user,
        start_time=start_time,
        tags=[RunTag(key, value) for key, value in tags.items()],
        run_name=run_name,
    )
    mock_store.reset_mock()
    MlflowClient().create_run(experiment_id, start_time, tags)
    mock_store.create_run.assert_called_once_with(
        experiment_id=experiment_id,
        user_id=user,
        start_time=start_time,
        tags=[RunTag(key, value) for key, value in tags.items()],
        run_name=None,
    )


def test_client_set_terminated_no_change_name(mock_store):
    experiment_id = mock.Mock()
    run = MlflowClient().create_run(experiment_id, run_name="my name")
    MlflowClient().set_terminated(run.info.run_id)
    _, kwargs = mock_store.update_run_info.call_args
    assert kwargs["run_name"] is None


def test_client_search_runs_defaults(mock_store):
    MlflowClient().search_runs([1, 2, 3])
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=[1, 2, 3],
        filter_string="",
        run_view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=None,
        page_token=None,
    )


def test_client_search_runs_filter(mock_store):
    MlflowClient().search_runs(["a", "b", "c"], "my filter")
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=["a", "b", "c"],
        filter_string="my filter",
        run_view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=None,
        page_token=None,
    )


def test_client_search_runs_view_type(mock_store):
    MlflowClient().search_runs(["a", "b", "c"], "my filter", ViewType.DELETED_ONLY)
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=["a", "b", "c"],
        filter_string="my filter",
        run_view_type=ViewType.DELETED_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=None,
        page_token=None,
    )


def test_client_search_runs_max_results(mock_store):
    MlflowClient().search_runs([5], "my filter", ViewType.ALL, 2876)
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=[5],
        filter_string="my filter",
        run_view_type=ViewType.ALL,
        max_results=2876,
        order_by=None,
        page_token=None,
    )


def test_client_search_runs_int_experiment_id(mock_store):
    MlflowClient().search_runs(123)
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=[123],
        filter_string="",
        run_view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=None,
        page_token=None,
    )


def test_client_search_runs_string_experiment_id(mock_store):
    MlflowClient().search_runs("abc")
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=["abc"],
        filter_string="",
        run_view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=None,
        page_token=None,
    )


def test_client_search_runs_order_by(mock_store):
    MlflowClient().search_runs([5], order_by=["a", "b"])
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=[5],
        filter_string="",
        run_view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=["a", "b"],
        page_token=None,
    )


def test_client_search_runs_page_token(mock_store):
    MlflowClient().search_runs([5], page_token="blah")
    mock_store.search_runs.assert_called_once_with(
        experiment_ids=[5],
        filter_string="",
        run_view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        order_by=None,
        page_token="blah",
    )


def test_update_registered_model(mock_registry_store):
    """
    Update registered model no longer supports name change.
    """
    expected_return_value = "some expected return value."
    mock_registry_store.rename_registered_model.return_value = expected_return_value
    expected_return_value_2 = "other expected return value."
    mock_registry_store.update_registered_model.return_value = expected_return_value_2
    res = MlflowClient(registry_uri="sqlite:///somedb.db").update_registered_model(
        name="orig name", description="new description"
    )
    assert expected_return_value_2 == res
    mock_registry_store.update_registered_model.assert_called_once_with(
        name="orig name", description="new description", deployment_job_id=None
    )
    mock_registry_store.rename_registered_model.assert_not_called()


def test_create_model_version(mock_registry_store):
    """
    Basic test for create model version.
    """
    mock_registry_store.create_model_version.return_value = _default_model_version()
    res = MlflowClient(registry_uri="sqlite:///somedb.db").create_model_version(
        "orig name", "source", "run-id", tags={"key": "value"}, description="desc"
    )
    assert res == _default_model_version()
    mock_registry_store.create_model_version.assert_called_once_with(
        "orig name",
        "source",
        "run-id",
        [ModelVersionTag(key="key", value="value")],
        None,
        "desc",
        local_model_path=None,
        model_id=None,
    )


def test_update_model_version(mock_registry_store):
    """
    Update registered model no longer support state changes.
    """
    mock_registry_store.update_model_version.return_value = _default_model_version()
    res = MlflowClient(registry_uri="sqlite:///somedb.db").update_model_version(
        name="orig name", version="1", description="desc"
    )
    assert _default_model_version() == res
    mock_registry_store.update_model_version.assert_called_once_with(
        name="orig name", version="1", description="desc"
    )
    mock_registry_store.transition_model_version_stage.assert_not_called()


def test_transition_model_version_stage(mock_registry_store):
    name = "Model 1"
    version = "12"
    stage = "Production"
    expected_result = ModelVersion(name, version, creation_timestamp=123, current_stage=stage)
    mock_registry_store.transition_model_version_stage.return_value = expected_result
    actual_result = MlflowClient(registry_uri="sqlite:///somedb.db").transition_model_version_stage(
        name, version, stage
    )
    mock_registry_store.transition_model_version_stage.assert_called_once_with(
        name=name, version=version, stage=stage, archive_existing_versions=False
    )
    assert expected_result == actual_result


def test_registry_uri_set_as_param():
    uri = "sqlite:///somedb.db"
    client = MlflowClient(tracking_uri="databricks://tracking", registry_uri=uri)
    assert client._registry_uri == uri


def test_registry_uri_from_set_registry_uri():
    uri = "sqlite:///somedb.db"
    set_registry_uri(uri)
    client = MlflowClient(tracking_uri="databricks://tracking")
    assert client._registry_uri == uri
    set_registry_uri(None)


def test_registry_uri_from_tracking_uri_param():
    tracking_uri = "databricks://tracking_vhawoierj"
    with mock.patch(
        "mlflow.tracking._tracking_service.utils.get_tracking_uri",
        return_value=tracking_uri,
    ):
        client = MlflowClient(tracking_uri=tracking_uri)
        # For databricks tracking URIs, registry URI defaults to Unity Catalog with profile
        assert client._registry_uri == "databricks-uc://tracking_vhawoierj"


def test_registry_uri_from_implicit_tracking_uri():
    tracking_uri = "databricks://tracking_wierojasdf"
    with mock.patch(
        "mlflow.tracking._tracking_service.utils.get_tracking_uri",
        return_value=tracking_uri,
    ):
        client = MlflowClient()
        # For databricks tracking URIs, registry URI defaults to Unity Catalog with profile
        assert client._registry_uri == "databricks-uc://tracking_wierojasdf"


def test_create_model_version_nondatabricks_source_no_runlink(mock_registry_store):
    run_id = "runid"
    client = MlflowClient(tracking_uri="http://10.123.1231.11")
    mock_registry_store.create_model_version.return_value = ModelVersion(
        "name",
        1,
        0,
        1,
        source="source",
        run_id=run_id,
    )
    model_version = client.create_model_version("name", "source", "runid")
    assert model_version.name == "name"
    assert model_version.source == "source"
    assert model_version.run_id == "runid"
    # verify that the store was not provided a run link
    mock_registry_store.create_model_version.assert_called_once_with(
        "name", "source", "runid", [], None, None, local_model_path=None, model_id=None
    )


def test_create_model_version_nondatabricks_source_no_run_id(mock_registry_store):
    client = MlflowClient(tracking_uri="http://10.123.1231.11")
    mock_registry_store.create_model_version.return_value = ModelVersion(
        "name", 1, 0, 1, source="source"
    )
    model_version = client.create_model_version("name", "source")
    assert model_version.name == "name"
    assert model_version.source == "source"
    assert model_version.run_id is None
    # verify that the store was not provided a run id
    mock_registry_store.create_model_version.assert_called_once_with(
        "name", "source", None, [], None, None, local_model_path=None, model_id=None
    )


def test_create_model_version_explicitly_set_run_link(
    mock_registry_store, mock_databricks_tracking_store
):
    run_link = "my-run-link"
    hostname = "https://workspace.databricks.com/"
    workspace_id = "10002"
    mock_registry_store.create_model_version.return_value = ModelVersion(
        "name",
        1,
        0,
        1,
        source="source",
        run_id=mock_databricks_tracking_store.run_id,
        run_link=run_link,
    )

    # mocks to make sure that even if you're in a notebook, this setting is respected.
    with (
        mock.patch("mlflow.utils.databricks_utils.is_in_databricks_notebook", return_value=True),
        mock.patch(
            "mlflow.utils.databricks_utils.get_workspace_info_from_dbutils",
            return_value=(hostname, workspace_id),
        ),
    ):
        client = MlflowClient(tracking_uri="databricks", registry_uri="otherplace")
        model_version = client.create_model_version("name", "source", "runid", run_link=run_link)
        assert model_version.run_link == run_link
        # verify that the store was provided with the explicitly passed in run link
        mock_registry_store.create_model_version.assert_called_once_with(
            "name", "source", "runid", [], run_link, None, local_model_path=None, model_id=None
        )


def test_create_model_version_run_link_in_notebook_with_default_profile(
    mock_registry_store, mock_databricks_tracking_store
):
    hostname = "https://workspace.databricks.com/"
    workspace_id = "10002"
    workspace_url = _construct_databricks_run_url(
        hostname,
        mock_databricks_tracking_store.experiment_id,
        mock_databricks_tracking_store.run_id,
        workspace_id,
    )

    with (
        mock.patch("mlflow.utils.databricks_utils.is_in_databricks_notebook", return_value=True),
        mock.patch(
            "mlflow.utils.databricks_utils.get_workspace_info_from_dbutils",
            return_value=(hostname, workspace_id),
        ),
    ):
        client = MlflowClient(tracking_uri="databricks", registry_uri="otherplace")
        mock_registry_store.create_model_version.return_value = ModelVersion(
            "name",
            1,
            0,
            1,
            source="source",
            run_id=mock_databricks_tracking_store.run_id,
            run_link=workspace_url,
        )
        model_version = client.create_model_version("name", "source", "runid")
        assert model_version.run_link == workspace_url
        # verify that the client generated the right URL
        mock_registry_store.create_model_version.assert_called_once_with(
            "name", "source", "runid", [], workspace_url, None, local_model_path=None, model_id=None
        )


def test_create_model_version_with_source(mock_registry_store, mock_databricks_tracking_store):
    model_id = "model_id"
    mock_registry_store.create_model_version.return_value = ModelVersion(
        "name",
        1,
        0,
        1,
        source="/path/to/source",
        run_id=mock_databricks_tracking_store.run_id,
        run_link=None,
        model_id=model_id,
    )
    mock_logged_model = LoggedModel(
        experiment_id="exp_id",
        model_id="model_id",
        name="name",
        artifact_location="/path/to/source",
        creation_timestamp=0,
        last_updated_timestamp=0,
    )

    with mock.patch(
        "mlflow.tracking.client.MlflowClient.get_logged_model", return_value=mock_logged_model
    ):
        client = MlflowClient(tracking_uri="databricks")
        model_version = client.create_model_version(
            "name", f"models:/{model_id}", "runid", run_link=None, model_id=model_id
        )
        assert model_version.model_id == model_id
        mock_registry_store.create_model_version.assert_called_once_with(
            "name",
            f"models:/{model_id}",
            "runid",
            [],
            None,
            None,
            local_model_path=None,
            model_id="model_id",
        )

    mock_registry_store.create_model_version.reset_mock()
    with mock.patch(
        "mlflow.tracking.client.MlflowClient.get_logged_model", return_value=mock_logged_model
    ):
        client = MlflowClient(tracking_uri="databricks", registry_uri="databricks-uc")
        model_version = client.create_model_version(
            "name", f"models:/{model_id}", "runid", run_link=None, model_id=model_id
        )
        assert model_version.model_id == model_id
        mock_registry_store.create_model_version.assert_called_once_with(
            "name",
            f"models:/{model_id}",
            "runid",
            [],
            None,
            None,
            local_model_path=None,
            model_id="model_id",
        )


def test_create_model_version_with_nondatabricks_source_uc_registry(mock_registry_store):
    name = "name"
    model_id = "model_id"
    run_id = "runid"
    source = "/path/to/source"
    model_uri = f"models:/{model_id}"
    mock_registry_store.create_model_version.return_value = ModelVersion(
        "name",
        1,
        0,
        1,
        source=source,
        run_id=run_id,
        run_link=None,
        model_id=model_id,
    )
    mock_logged_model = LoggedModel(
        experiment_id="exp_id",
        model_id=model_id,
        name=name,
        artifact_location=source,
        creation_timestamp=0,
        last_updated_timestamp=0,
    )

    with mock.patch(
        "mlflow.tracking.client.MlflowClient.get_logged_model", return_value=mock_logged_model
    ):
        client = MlflowClient(tracking_uri="http://10.123.1231.11", registry_uri="databricks-uc")
        model_version = client.create_model_version(
            name, model_uri, run_id, run_link=None, model_id=model_id
        )
        assert model_version.model_id == model_id
        mock_registry_store.create_model_version.assert_called_once_with(
            name,
            source,
            run_id,
            [],
            None,
            None,
            local_model_path=None,
            model_id=None,
        )


def test_creation_default_values_in_unity_catalog(mock_registry_store):
    client = MlflowClient(tracking_uri="databricks", registry_uri="databricks-uc")
    mock_registry_store.create_model_version.return_value = ModelVersion(
        "name",
        1,
        0,
        1,
        source="source",
        run_id="runid",
    )
    client.create_model_version("name", "source", "runid")
    # verify that registry store was called with tags=[] and run_link=None
    mock_registry_store.create_model_version.assert_called_once_with(
        "name", "source", "runid", [], None, None, local_model_path=None, model_id=None
    )
    client.create_registered_model(name="name", description="description")
    # verify that registry store was called with tags=[]
    mock_registry_store.create_registered_model.assert_called_once_with(
        "name", [], "description", None
    )


def test_await_model_version_creation(mock_registry_store):
    mv = ModelVersion(
        name="name",
        version=1,
        creation_timestamp=123,
        status=ModelVersionStatus.to_string(ModelVersionStatus.FAILED_REGISTRATION),
    )
    mock_registry_store.create_model_version.return_value = mv

    client = MlflowClient(tracking_uri="http://10.123.1231.11")

    client.create_model_version("name", "source")
    mock_registry_store._await_model_version_creation.assert_called_once_with(
        mv, DEFAULT_AWAIT_MAX_SLEEP_SECONDS
    )


def test_create_model_version_run_link_with_configured_profile(
    mock_registry_store, mock_databricks_tracking_store
):
    hostname = "https://workspace.databricks.com/"
    workspace_id = "10002"
    workspace_url = _construct_databricks_run_url(
        hostname,
        mock_databricks_tracking_store.experiment_id,
        mock_databricks_tracking_store.run_id,
        workspace_id,
    )

    with (
        mock.patch("mlflow.utils.databricks_utils.is_in_databricks_notebook", return_value=False),
        mock.patch(
            "mlflow.utils.databricks_utils.get_workspace_info_from_databricks_secrets",
            return_value=(hostname, workspace_id),
        ),
    ):
        client = MlflowClient(tracking_uri="databricks", registry_uri="otherplace")
        mock_registry_store.create_model_version.return_value = ModelVersion(
            "name",
            1,
            0,
            1,
            source="source",
            run_id=mock_databricks_tracking_store.run_id,
            run_link=workspace_url,
        )
        model_version = client.create_model_version("name", "source", "runid")
        assert model_version.run_link == workspace_url
        # verify that the client generated the right URL
        mock_registry_store.create_model_version.assert_called_once_with(
            "name", "source", "runid", [], workspace_url, None, local_model_path=None, model_id=None
        )


def test_create_model_version_copy_called_db_to_db(mock_registry_store):
    client = MlflowClient(
        tracking_uri="databricks://tracking",
        registry_uri="databricks://registry:workspace",
    )
    mock_registry_store.create_model_version.return_value = _default_model_version()
    with mock.patch("mlflow.tracking.client._upload_artifacts_to_databricks") as upload_mock:
        client.create_model_version(
            "model name",
            "dbfs:/source",
            "run_12345",
            run_link="not:/important/for/test",
        )
        upload_mock.assert_called_once_with(
            "dbfs:/source",
            "run_12345",
            "databricks://tracking",
            "databricks://registry:workspace",
        )


def test_create_model_version_copy_called_nondb_to_db(mock_registry_store):
    client = MlflowClient(
        tracking_uri="https://tracking", registry_uri="databricks://registry:workspace"
    )
    mock_registry_store.create_model_version.return_value = _default_model_version()
    with mock.patch("mlflow.tracking.client._upload_artifacts_to_databricks") as upload_mock:
        client.create_model_version(
            "model name", "s3:/source", "run_12345", run_link="not:/important/for/test"
        )
        upload_mock.assert_called_once_with(
            "s3:/source",
            "run_12345",
            "https://tracking",
            "databricks://registry:workspace",
        )


def test_create_model_version_copy_not_called_to_db(mock_registry_store):
    client = MlflowClient(
        tracking_uri="databricks://registry:workspace",
        registry_uri="databricks://registry:workspace",
    )
    mock_registry_store.create_model_version.return_value = _default_model_version()
    with mock.patch("mlflow.tracking.client._upload_artifacts_to_databricks") as upload_mock:
        client.create_model_version(
            "model name",
            "dbfs:/source",
            "run_12345",
            run_link="not:/important/for/test",
        )
        upload_mock.assert_not_called()


def test_create_model_version_copy_not_called_to_nondb(mock_registry_store):
    client = MlflowClient(tracking_uri="databricks://tracking", registry_uri="https://registry")
    mock_registry_store.create_model_version.return_value = _default_model_version()
    with mock.patch("mlflow.tracking.client._upload_artifacts_to_databricks") as upload_mock:
        client.create_model_version(
            "model name",
            "dbfs:/source",
            "run_12345",
            run_link="not:/important/for/test",
        )
        upload_mock.assert_not_called()


def _default_model_version():
    return ModelVersion("model name", 1, creation_timestamp=123, status="READY")


def test_client_can_be_serialized_with_pickle(tmp_path):
    """
    Verifies that instances of `MlflowClient` can be serialized using pickle, even if the underlying
    Tracking and Model Registry stores used by the client are not serializable using pickle
    """

    class MockUnpickleableTrackingStore(SqlAlchemyTrackingStore):
        pass

    class MockUnpickleableModelRegistryStore(SqlAlchemyModelRegistryStore):
        pass

    backend_store_path = tmp_path.joinpath("test.db")
    artifact_store_path = tmp_path.joinpath("artifacts")

    mock_tracking_store = MockUnpickleableTrackingStore(
        f"sqlite:///{backend_store_path}", str(artifact_store_path)
    )
    mock_model_registry_store = MockUnpickleableModelRegistryStore(
        f"sqlite:///{backend_store_path}"
    )

    # Verify that the mock stores cannot be pickled because they are defined within a function
    # (i.e. the test function)
    with pytest.raises(AttributeError, match="<locals>.MockUnpickleableTrackingStore'"):
        pickle.dumps(mock_tracking_store)

    with pytest.raises(AttributeError, match="<locals>.MockUnpickleableModelRegistryStore'"):
        pickle.dumps(mock_model_registry_store)

    _register("pickle", lambda *args, **kwargs: mock_tracking_store)
    _get_model_registry_store_registry().register(
        "pickle", lambda *args, **kwargs: mock_model_registry_store
    )

    # Create an MlflowClient with the store that cannot be pickled, perform
    # tracking & model registry operations, and verify that the client can still be pickled
    client = MlflowClient("pickle://foo")
    client.create_experiment("test_experiment")
    client.create_registered_model("test_model")
    pickle.dumps(client)


@pytest.fixture
def mock_registry_store_with_get_latest_version(mock_registry_store):
    mock_get_latest_versions = mock.Mock()
    mock_get_latest_versions.return_value = [
        ModelVersion(
            "model_name",
            1,
            0,
        )
    ]

    mock_registry_store.get_latest_versions = mock_get_latest_versions
    return mock_registry_store


def test_set_model_version_tag(mock_registry_store_with_get_latest_version):
    # set_model_version_tag using version
    MlflowClient().set_model_version_tag("model_name", 1, "tag1", "foobar")
    mock_registry_store_with_get_latest_version.set_model_version_tag.assert_called_once_with(
        "model_name", 1, ModelVersionTag(key="tag1", value="foobar")
    )

    mock_registry_store_with_get_latest_version.set_model_version_tag.reset_mock()

    # set_model_version_tag using stage
    MlflowClient().set_model_version_tag("model_name", key="tag1", value="foobar", stage="Staging")
    mock_registry_store_with_get_latest_version.set_model_version_tag.assert_called_once_with(
        "model_name", 1, ModelVersionTag(key="tag1", value="foobar")
    )

    # set_model_version_tag with version and stage set
    with pytest.raises(MlflowException, match="version and stage cannot be set together"):
        MlflowClient().set_model_version_tag("model_name", 1, "tag1", "foobar", stage="Staging")

    # set_model_version_tag with version and stage not set
    with pytest.raises(MlflowException, match="version or stage must be set"):
        MlflowClient().set_model_version_tag("model_name", key="tag1", value="foobar")


def test_delete_model_version_tag(mock_registry_store_with_get_latest_version):
    # delete_model_version_tag using version
    MlflowClient().delete_model_version_tag("model_name", 1, "tag1")
    mock_registry_store_with_get_latest_version.delete_model_version_tag.assert_called_once_with(
        "model_name", 1, "tag1"
    )

    mock_registry_store_with_get_latest_version.delete_model_version_tag.reset_mock()

    # delete_model_version_tag using stage
    MlflowClient().delete_model_version_tag("model_name", key="tag1", stage="Staging")
    mock_registry_store_with_get_latest_version.delete_model_version_tag.assert_called_once_with(
        "model_name", 1, "tag1"
    )

    # delete_model_version_tag with version and stage set
    with pytest.raises(MlflowException, match="version and stage cannot be set together"):
        MlflowClient().delete_model_version_tag(
            "model_name", version=1, key="tag1", stage="staging"
        )

    # delete_model_version_tag with version and stage not set
    with pytest.raises(MlflowException, match="version or stage must be set"):
        MlflowClient().delete_model_version_tag("model_name", key="tag1")


def test_set_registered_model_alias(mock_registry_store):
    MlflowClient().set_registered_model_alias("model_name", "test_alias", 1)
    mock_registry_store.set_registered_model_alias.assert_called_once_with(
        "model_name", "test_alias", 1
    )


def test_delete_registered_model_alias(mock_registry_store):
    MlflowClient().delete_registered_model_alias("model_name", "test_alias")
    mock_registry_store.delete_registered_model_alias.assert_called_once_with(
        "model_name", "test_alias"
    )


def test_get_model_version_by_alias(mock_registry_store):
    mock_registry_store.get_model_version_by_alias.return_value = _default_model_version()
    res = MlflowClient().get_model_version_by_alias("model_name", "test_alias")
    assert res == _default_model_version()
    mock_registry_store.get_model_version_by_alias.assert_called_once_with(
        "model_name", "test_alias"
    )


def test_update_run(mock_store):
    MlflowClient().update_run(run_id="run_id", status="FINISHED", name="my name")
    mock_store.update_run_info.assert_called_once_with(
        run_id="run_id",
        run_status=RunStatus.from_string("FINISHED"),
        end_time=mock.ANY,
        run_name="my name",
    )


def test_client_log_metric_params_tags_overrides(mock_store):
    experiment_id = mock.Mock()
    start_time = mock.Mock()
    run_name = mock.Mock()
    run = MlflowClient().create_run(experiment_id, start_time, tags={}, run_name=run_name)
    run_id = run.info.run_id

    run_operation = MlflowClient().log_metric(run_id, "m1", 0.87, 123456789, 1, synchronous=False)
    run_operation.wait()

    run_operation = MlflowClient().log_param(run_id, "p1", "pv1", synchronous=False)
    run_operation.wait()

    run_operation = MlflowClient().set_tag(run_id, "t1", "tv1", synchronous=False)
    run_operation.wait()

    mock_store.log_metric_async.assert_called_once_with(run_id, Metric("m1", 0.87, 123456789, 1))
    mock_store.log_param_async.assert_called_once_with(run_id, Param("p1", "pv1"))
    mock_store.set_tag_async.assert_called_once_with(run_id, RunTag("t1", "tv1"))

    mock_store.reset_mock()

    # log_batch_async
    MlflowClient().create_run(experiment_id, start_time, {})
    metrics = [Metric("m1", 0.87, 123456789, 1), Metric("m2", 0.87, 123456789, 1)]
    tags = [RunTag("t1", "tv1"), RunTag("t2", "tv2")]
    params = [Param("p1", "pv1"), Param("p2", "pv2")]
    run_operation = MlflowClient().log_batch(run_id, metrics, params, tags, synchronous=False)
    run_operation.wait()

    mock_store.log_batch_async.assert_called_once_with(
        run_id=run_id, metrics=metrics, params=params, tags=tags
    )


def test_invalid_run_id_log_artifact():
    with pytest.raises(
        MlflowException,
        match=r"Invalid run id.*",
    ):
        MlflowClient().log_artifact("tr-123", "path")


def test_enable_async_logging(mock_store, setup_async_logging):
    MlflowClient().log_param(run_id="run_id", key="key", value="val")
    mock_store.log_param_async.assert_called_once_with("run_id", Param("key", "val"))

    MlflowClient().log_metric(run_id="run_id", key="key", value="val", step=1, timestamp=1)
    mock_store.log_metric_async.assert_called_once_with("run_id", Metric("key", "val", 1, 1))


def test_file_store_download_upload_trace_data(tmp_path):
    pytest.skip("FileStore is no longer supported.")
    with _use_tracking_uri(tmp_path.joinpath("mlruns").as_uri()):
        client = MlflowClient()
        span = client.start_trace("test", inputs={"test": 1})
        client.end_trace(span.trace_id, outputs={"result": 2})
        trace = mlflow.get_trace(span.trace_id, flush=True)
        trace_data = client.get_trace(span.trace_id, flush=True).data
        assert trace_data.request == trace.data.request
        assert trace_data.response == trace.data.response


def test_get_trace_throw_if_trace_id_is_online_trace_id(db_uri):
    client = MlflowClient("databricks")
    trace_id = "3a3c3b56-910a-4721-8d02-0333eda5f37e"
    with pytest.raises(MlflowException, match="Traces from inference tables can only be loaded"):
        client.get_trace(trace_id)

    another_client = MlflowClient(db_uri)
    with pytest.raises(MlflowException, match=r"Trace with ID '[\w-]+' not found"):
        another_client.get_trace(trace_id)


@pytest.fixture(params=["file", "sqlalchemy"])
def registry_uri(request, tmp_path, db_uri):
    """Set an MLflow Model Registry URI with different type of backend."""
    if request.param == "file":
        pytest.skip("FileStore is no longer supported.")
    if "MLFLOW_SKINNY" in os.environ and request.param == "sqlalchemy":
        pytest.skip("SQLAlchemy store is not available in skinny.")

    original_registry_uri = mlflow.get_registry_uri()

    if request.param == "file":
        registry_uri = tmp_path.joinpath("file").as_uri()
    elif request.param == "sqlalchemy":
        registry_uri = db_uri

    yield registry_uri

    # Reset tracking URI
    mlflow.set_tracking_uri(original_registry_uri)


def test_crud_prompts(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    client.register_prompt(
        name="prompt_1",
        template="Hi, {{title}} {{name}}! How are you today?",
        commit_message="A friendly greeting",
    )

    prompt = client.load_prompt("prompt_1", version=1)
    assert prompt.name == "prompt_1"
    assert prompt.template == "Hi, {{title}} {{name}}! How are you today?"
    assert prompt.commit_message == "A friendly greeting"

    client.register_prompt(
        name="prompt_1",
        template="Hi, {{title}} {{name}}! What's up?",
        commit_message="New greeting",
    )

    prompt = client.load_prompt("prompt_1", version=2)
    assert prompt.template == "Hi, {{title}} {{name}}! What's up?"

    prompt = client.load_prompt("prompt_1", version=1)
    assert prompt.template == "Hi, {{title}} {{name}}! How are you today?"

    prompt = client.load_prompt("prompts:/prompt_1/2")
    assert prompt.template == "Hi, {{title}} {{name}}! What's up?"

    # Test loading non-existent prompts
    assert mlflow.load_prompt("does_not_exist", version=1, allow_missing=True) is None


def test_create_prompt_with_tags_and_metadata(tracking_uri, disable_prompt_cache):
    def wait_for_prompt_linking():
        """Wait for background prompt linking threads to complete."""
        for t in threading.enumerate():
            if t.name.startswith("link_prompt_to_experiment_thread"):
                t.join(timeout=5.0)
                if t.is_alive():
                    raise TimeoutError(f"Thread {t.name} did not complete within timeout.")

    client = MlflowClient(tracking_uri=tracking_uri)

    # Create prompt with version-specific tags
    client.register_prompt(
        name="prompt_1",
        template="Hi, {{name}}!",
        tags={"author": "Alice"},  # This will be version-level tags now
    )

    # Wait for the background linking thread to complete
    wait_for_prompt_linking()

    # Set some prompt-level tags separately
    client.set_prompt_tag("prompt_1", "application", "greeting")
    client.set_prompt_tag("prompt_1", "language", "en")

    # Test version 1
    prompt_v1 = client.load_prompt("prompt_1", version=1)
    assert prompt_v1.template == "Hi, {{name}}!"
    # Version tags are separate from prompt tags
    assert prompt_v1.tags == {"author": "Alice"}

    # Wait for the background linking thread from load_prompt
    wait_for_prompt_linking()

    # Test prompt-level tags (separate from version)
    prompt_entity = client.get_prompt("prompt_1")
    # Note: Currently includes the version tags too, but we expect this behavior to change
    assert prompt_entity.tags == {
        "author": "Alice",  # This appears due to current implementation
        "application": "greeting",
        "language": "en",
        "_mlflow_experiment_ids": ",0,",  # Linked to Default experiment
    }

    # Create version 2 with different version-level tags
    client.register_prompt(
        name="prompt_1",
        template="こんにちは、{{name}}!",
        tags={"author": "Bob", "date": "2022-01-01"},  # Version-level tags
    )

    # Wait for the background linking thread from register_prompt
    wait_for_prompt_linking()

    # Update some prompt-level tags
    client.set_prompt_tag("prompt_1", "project", "toy")
    client.set_prompt_tag("prompt_1", "language", "ja")

    # Test version 2
    prompt_v2 = client.load_prompt("prompt_1", version=2)
    assert prompt_v2.template == "こんにちは、{{name}}!"
    # Version 2 has its own version tags (decoupled from prompt and version 1)
    assert prompt_v2.tags == {"author": "Bob", "date": "2022-01-01"}

    # Wait for the background linking thread from load_prompt
    wait_for_prompt_linking()

    # Verify prompt-level tags are updated and separate
    prompt_entity_updated = client.get_prompt("prompt_1")
    # Note: Currently the prompt tags get overwritten by the newest version's tags
    assert prompt_entity_updated.tags == {
        "author": "Bob",  # This appears due to current implementation
        "date": "2022-01-01",  # This appears due to current implementation
        "application": "greeting",
        "project": "toy",
        "language": "ja",
        "_mlflow_experiment_ids": ",0,",  # Linked to Default experiment
    }

    # Version 1 tags should be unchanged (decoupled from prompt tags)
    prompt_v1_after_update = client.load_prompt("prompt_1", version=1)
    assert prompt_v1_after_update.tags == {"author": "Alice"}  # Unchanged


def test_create_prompt_error_handling(tracking_uri, disable_prompt_cache):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Exceeds the max length
    with pytest.raises(MlflowException, match=r"Prompt text exceeds max length of"):
        client.register_prompt(name="prompt_1", template="a" * 100_001)

    # When the first version creation fails, RegisteredModel should not be created
    with pytest.raises(MlflowException, match=r"Prompt with name=prompt_1 not found"):
        client.load_prompt("prompt_1", version=1)

    client.register_prompt("prompt_1", template="Hi, {{title}} {{name}}!")
    assert client.load_prompt("prompt_1", version=1) is not None

    # When the subsequent version creation fails, RegisteredModel should remain
    with pytest.raises(MlflowException, match=r"Prompt text exceeds max length of"):
        client.register_prompt(name="prompt_1", template="a" * 100_001)

    assert client.load_prompt("prompt_1", version=1) is not None


def test_create_prompt_with_invalid_name(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    with pytest.raises(MlflowException, match=r"Prompt name must be a non-empty string"):
        client.register_prompt(name="", template="Hi, {{name}}!")

    with pytest.raises(MlflowException, match=r"Prompt name must be a non-empty string"):
        client.register_prompt(name=123, template="Hi, {{name}}!")

    for invalid_pattern in [
        "prompt_1/2",
        "m%6fdel",
        "prompt?!?",
        "prompt with space",
    ]:
        with pytest.raises(MlflowException, match=r"Prompt name can only contain alphanumeric"):
            client.register_prompt(name=invalid_pattern, template="Hi, {{name}}!")

    # Name conflicts with a model
    client.create_registered_model("model")
    with pytest.raises(MlflowException, match=r"Model 'model' exists with the same name."):
        client.register_prompt(name="model", template="Hi, {{name}}!")


def test_load_prompt_error(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    with pytest.raises(MlflowException, match=r"Prompt with name=test not found"):
        client.load_prompt("test", version=1)

    # Both file and sqlalchemy return the same error format now
    error_msg = r"Prompt with name=test not found"

    with pytest.raises(MlflowException, match=error_msg):
        client.load_prompt("test", version=2)

    with pytest.raises(MlflowException, match=error_msg):
        client.load_prompt("test", version=2, allow_missing=False)

    # Load prompt with a model name
    client.create_registered_model("model")
    client.create_model_version("model", "source")

    with pytest.raises(MlflowException, match=r"Name `model` is registered as a model"):
        client.load_prompt("model", version=1)

    with pytest.raises(MlflowException, match=r"Name `model` is registered as a model"):
        client.load_prompt("model", version=1)

    with pytest.raises(MlflowException, match=r"Name `model` is registered as a model"):
        client.load_prompt("model", version=1, allow_missing=False)

    with pytest.raises(MlflowException, match=r"Name `model` is registered as a model"):
        client.load_prompt("model", version=1, allow_missing=False)


def test_link_prompt_version_to_run(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    prompt = client.register_prompt("prompt", template="Hi, {{name}}!")

    # Create actual runs to link to
    run1 = client.create_run(experiment_id="0").info.run_id
    run2 = client.create_run(experiment_id="0").info.run_id

    # Test that the method can be called without error
    client.link_prompt_version_to_run(run1, prompt)
    client.link_prompt_version_to_run(run2, prompt)

    # Verify tag was set by checking the run data
    run_data = client.get_run(run1)
    linked_prompts_tag = run_data.data.tags.get("mlflow.linkedPrompts")
    assert linked_prompts_tag is not None

    # Verify the JSON structure
    linked_prompts = json.loads(linked_prompts_tag)
    assert any(p["name"] == "prompt" and p["version"] == "1" for p in linked_prompts)

    # Test error case
    with pytest.raises(MlflowException, match=r"The `prompt` argument must be"):
        client.link_prompt_version_to_run(run1, 123)


@pytest.mark.parametrize("registry_uri", ["databricks"])
def test_crud_prompt_on_unsupported_registry(registry_uri):
    client = MlflowClient(registry_uri=registry_uri)

    with pytest.raises(MlflowException, match=r"The 'register_prompt' API is not supported"):
        client.register_prompt(
            name="prompt_1",
            template="Hi, {{title}} {{name}}! How are you today?",
            commit_message="A friendly greeting",
            tags={"model": "my-model"},
        )

    with pytest.raises(MlflowException, match=r"The 'load_prompt' API is not supported"):
        client.load_prompt("prompt_1")


def test_block_create_model_with_prompt_tag(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    with pytest.raises(MlflowException, match=r"Prompts cannot be registered"):
        client.create_registered_model(
            name="model",
            tags={IS_PROMPT_TAG_KEY: "true"},
        )

    with pytest.raises(MlflowException, match=r"Prompts cannot be registered"):
        client.create_model_version(
            name="model",
            source="source",
            tags={IS_PROMPT_TAG_KEY: "false"},
        )


def test_block_create_prompt_with_existing_model_name(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    client.create_registered_model("model")

    with pytest.raises(MlflowException, match=r"Model 'model' exists with"):
        client.register_prompt(
            name="model",
            template="Hi, {{title}} {{name}}! How are you today?",
            commit_message="A friendly greeting",
            tags={"model": "my-model"},
        )


def test_block_handling_prompt_with_model_apis(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)
    client.register_prompt("prompt", template="Hi, {{name}}!")
    client.set_prompt_alias("prompt", alias="alias", version=1)
    # Validate the prompt is registered
    prompt = client.load_prompt("prompt", version=1)
    assert prompt.name == "prompt"
    assert prompt.aliases == ["alias"]

    apis_to_args = [
        (client.rename_registered_model, ["prompt", "new_name"]),
        (client.update_registered_model, ["prompt", "new_description"]),
        (client.delete_registered_model, ["prompt"]),
        (client.get_registered_model, ["prompt"]),
        (client.get_latest_versions, ["prompt"]),
        (client.set_registered_model_tag, ["prompt", "tag", "value"]),
        (client.delete_registered_model_tag, ["prompt", "tag"]),
        (client.update_model_version, ["prompt", 1, "new_description"]),
        (client.transition_model_version_stage, ["prompt", 1, "Production"]),
        (client.delete_model_version, ["prompt", 1]),
        (client.get_model_version, ["prompt", 1]),
        (client.get_model_version_download_uri, ["prompt", 1]),
        (client.set_model_version_tag, ["prompt", 1, "tag", "value"]),
        (client.delete_model_version_tag, ["prompt", 1, "tag"]),
        (client.set_registered_model_alias, ["prompt", "alias", 1]),
        (client.delete_registered_model_alias, ["prompt", "alias"]),
        (client.get_model_version_by_alias, ["prompt", "alias"]),
    ]

    for api, args in apis_to_args:
        with pytest.raises(MlflowException, match=r"Registered Model with name='prompt' not found"):
            api(*args)

    with pytest.raises(MlflowException, match=r"Model with uri 'models:/prompt/1' not found"):
        client.copy_model_version("models:/prompt/1", "new_model")


def test_log_and_detach_prompt(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    client.register_prompt(name="p1", template="Hi, there!")
    time.sleep(0.001)  # To avoid timestamp precision issue in Windows
    client.register_prompt(name="p2", template="Hi, {{name}}!")

    run_id = client.create_run(experiment_id="0").info.run_id

    # Check that initially no prompts are linked to the run
    run = client.get_run(run_id)
    linked_prompts_tag = run.data.tags.get(TraceTagKey.LINKED_PROMPTS)
    assert linked_prompts_tag is None

    client.link_prompt_version_to_run(run_id, "prompts:/p1/1")
    run = client.get_run(run_id)
    linked_prompts_tag = run.data.tags.get(TraceTagKey.LINKED_PROMPTS)
    assert linked_prompts_tag is not None
    prompts = json.loads(linked_prompts_tag)
    assert len(prompts) == 1
    assert prompts[0]["name"] == "p1"

    client.link_prompt_version_to_run(run_id, "prompts:/p2/1")
    run = client.get_run(run_id)
    linked_prompts_tag = run.data.tags.get(TraceTagKey.LINKED_PROMPTS)
    prompts = json.loads(linked_prompts_tag)
    assert len(prompts) == 2
    prompt_names = [p["name"] for p in prompts]
    assert "p1" in prompt_names
    assert "p2" in prompt_names


def test_search_prompt(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    client.register_prompt(name="prompt_1", template="Hi, {{name}}!")
    client.register_prompt(name="prompt_2", template="Hello, {{name}}!")
    client.register_prompt(name="prompt_3", template="Greetings, {{name}}!")
    client.register_prompt(name="prompt_4", template="Howdy, {{name}}!")
    client.register_prompt(name="prompt_5", template="Salutations, {{name}}!")
    client.register_prompt(name="prompt_6", template="Bonjour, {{name}}!")
    client.register_prompt(name="test", template="Test Template")
    client.register_prompt(name="new", template="Bonjour, {{name}}!")

    prompts = client.search_prompts(filter_string="name='prompt_1'")
    assert len(prompts) == 1
    assert prompts[0].name == "prompt_1"

    prompts = client.search_prompts(filter_string="name LIKE '%prompt%'")
    assert len(prompts) == 6
    assert all("prompt" in prompt.name for prompt in prompts)

    prompts = client.search_prompts()
    assert len(prompts) == 8

    prompts = client.search_prompts(max_results=3)
    assert len(prompts) == 3


def test_delete_prompt_version_no_auto_cleanup(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Create prompt and version
    client.register_prompt(name="test_prompt", template="Hello {{name}}!")

    # Verify prompt and version exist
    prompt = client.get_prompt("test_prompt")
    assert prompt is not None
    assert prompt.name == "test_prompt"

    prompt_version = client.get_prompt_version("test_prompt", 1)
    assert prompt_version is not None
    assert prompt_version.version == 1

    # Delete the version - prompt should remain
    client.delete_prompt_version("test_prompt", "1")

    # Prompt should still exist even though it has no versions
    prompt = client.get_prompt("test_prompt")
    assert prompt is not None
    assert prompt.name == "test_prompt"

    # Version should be gone
    with pytest.raises(MlflowException, match=r"Prompt.*name=test_prompt.*version=1.*not found"):
        client.get_prompt_version("test_prompt", 1)


def test_delete_prompt_version_invalidates_cached_load_prompt(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    prompt_ver = client.register_prompt(name="test_prompt", template="Version 1")
    loaded = client.load_prompt(prompt_ver.name, version=prompt_ver.version)
    assert loaded.template == "Version 1"

    client.delete_prompt_version(prompt_ver.name, str(prompt_ver.version))

    with pytest.raises(
        MlflowException,
        match=rf"Prompt.*name={prompt_ver.name}.*version={prompt_ver.version}.*not found",
    ):
        client.get_prompt_version(prompt_ver.name, prompt_ver.version)

    with pytest.raises(
        MlflowException,
        match=rf"Prompt.*name={prompt_ver.name}.*version={prompt_ver.version}.*not found",
    ):
        client.load_prompt(prompt_ver.name, version=prompt_ver.version)


def test_delete_prompt_version_invalidates_latest_cache(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    prompt_v1 = client.register_prompt(name="test_prompt", template="Version 1")
    prompt_v2 = client.register_prompt(name=prompt_v1.name, template="Version 2")

    latest_prompt = client.load_prompt(f"prompts:/{prompt_v1.name}@latest")
    assert latest_prompt.version == prompt_v2.version
    assert latest_prompt.template == prompt_v2.template

    client.delete_prompt_version(prompt_v2.name, str(prompt_v2.version))

    latest_prompt_after_delete = client.load_prompt(f"prompts:/{prompt_v1.name}@latest")
    assert latest_prompt_after_delete.version == prompt_v1.version
    assert latest_prompt_after_delete.template == prompt_v1.template


def test_set_prompt_model_config_invalidates_latest_cache(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    cache_ttl_seconds = 60
    prompt = client.register_prompt(name="test_prompt", template="test")
    prompt_before_update = client.load_prompt(prompt.name, cache_ttl_seconds=cache_ttl_seconds)
    assert prompt_before_update.model_config is None

    model_config = {"model_name": "gpt-4", "temperature": 0.7}
    mlflow.genai.set_prompt_model_config(
        name=prompt.name,
        version=prompt.version,
        model_config=model_config,
    )

    prompt_after_update = client.load_prompt(prompt.name, cache_ttl_seconds=cache_ttl_seconds)
    assert prompt_after_update.model_config == model_config


def test_delete_prompt_model_config_invalidates_latest_cache(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    cache_ttl_seconds = 60
    model_config = {"model_name": "gpt-4", "temperature": 0.7}
    prompt = client.register_prompt(
        name="test_prompt",
        template="test",
        model_config=model_config,
    )
    prompt_before_delete = client.load_prompt(prompt.name, cache_ttl_seconds=cache_ttl_seconds)
    assert prompt_before_delete.model_config == model_config

    mlflow.genai.delete_prompt_model_config(name=prompt.name, version=prompt.version)

    prompt_after_delete = client.load_prompt(prompt.name, cache_ttl_seconds=cache_ttl_seconds)
    assert prompt_after_delete.model_config is None


def test_delete_prompt_version_invalidates_alias_cache(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    prompt_v1 = client.register_prompt(name="test_prompt", template="Version 1")
    client.register_prompt(name=prompt_v1.name, template="Version 2")
    client.set_prompt_alias(prompt_v1.name, alias="production", version=prompt_v1.version)

    aliased_prompt = client.load_prompt(f"prompts:/{prompt_v1.name}@production")
    assert aliased_prompt.version == prompt_v1.version
    assert aliased_prompt.template == prompt_v1.template

    client.delete_prompt_version(prompt_v1.name, str(prompt_v1.version))

    with pytest.raises(
        MlflowException,
        match=(
            r"Prompt (.*) does not exist.|Prompt alias (.*) not found.|"
            rf"Prompt.*version={prompt_v1.version}.*not found"
        ),
    ):
        client.load_prompt(f"prompts:/{prompt_v1.name}@production")


def test_delete_prompt_with_no_versions(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)
    mlflow.set_experiment("test_delete_prompt_with_no_versions")

    # Create prompt and version, then delete version
    client.register_prompt(name="empty_prompt", template="Hello {{name}}!")
    client.delete_prompt_version("empty_prompt", "1")

    # Verify prompt exists but has no versions
    prompt = client.get_prompt("empty_prompt")
    assert prompt is not None

    # Delete the prompt - should work regardless of registry type
    client.delete_prompt("empty_prompt")

    # Prompt should be gone
    prompt = client.get_prompt("empty_prompt")
    assert prompt is None


def test_delete_prompt_invalidates_cached_load_prompt(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    prompt_ver = client.register_prompt(name="test_prompt", template="Version 1")
    loaded = client.load_prompt(prompt_ver.name, version=prompt_ver.version)
    assert loaded.template == "Version 1"

    client.delete_prompt(prompt_ver.name)

    assert client.get_prompt(prompt_ver.name) is None

    with pytest.raises(MlflowException, match=rf"Prompt.*name={prompt_ver.name}.*not found"):
        client.load_prompt(prompt_ver.name, version=prompt_ver.version)


def test_delete_prompt_complete_workflow(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Create prompt with multiple versions
    client.register_prompt(name="workflow_prompt", template="Version 1: {{name}}")
    client.register_prompt(name="workflow_prompt", template="Version 2: {{name}}")
    client.register_prompt(name="workflow_prompt", template="Version 3: {{name}}")

    # Verify all versions exist
    v1 = client.get_prompt_version("workflow_prompt", 1)
    v2 = client.get_prompt_version("workflow_prompt", 2)
    v3 = client.get_prompt_version("workflow_prompt", 3)
    assert v1.template == "Version 1: {{name}}"
    assert v2.template == "Version 2: {{name}}"
    assert v3.template == "Version 3: {{name}}"

    # Delete versions one by one
    client.delete_prompt_version("workflow_prompt", "1")
    client.delete_prompt_version("workflow_prompt", "2")
    client.delete_prompt_version("workflow_prompt", "3")

    # Prompt should still exist with no versions
    prompt = client.get_prompt("workflow_prompt")
    assert prompt is not None

    # Now delete the prompt itself
    client.delete_prompt("workflow_prompt")

    # Prompt should be completely gone
    prompt = client.get_prompt("workflow_prompt")
    assert prompt is None


def test_delete_prompt_error_handling(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Test deleting non-existent prompt
    with pytest.raises(MlflowException, match=r"Prompt with name=nonexistent not found"):
        client.delete_prompt("nonexistent")

    # Test deleting non-existent version
    client.register_prompt(name="test_errors", template="Hello {{name}}!")
    with pytest.raises(MlflowException, match=r"Prompt.*name=test_errors.*version=999.*not found"):
        client.delete_prompt_version("test_errors", "999")


def test_delete_prompt_version_behavior_consistency(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Create multiple prompts with versions
    for i in range(3):
        prompt_name = f"consistency_test_{i}"
        client.register_prompt(name=prompt_name, template=f"Template {i}: {{{{name}}}}")

        # Delete the version immediately
        client.delete_prompt_version(prompt_name, "1")

        # Prompt should remain but have no versions
        prompt = client.get_prompt(prompt_name)
        assert prompt is not None
        assert prompt.name == prompt_name

        # Version should be gone
        with pytest.raises(MlflowException, match=r"Prompt.*version.*not found"):
            client.get_prompt_version(prompt_name, 1)

    # Clean up - delete all prompts
    for i in range(3):
        client.delete_prompt(f"consistency_test_{i}")
        prompt = client.get_prompt(f"consistency_test_{i}")
        assert prompt is None


@pytest.mark.parametrize("registry_uri", ["databricks-uc"])
def test_delete_prompt_with_versions_unity_catalog_error(registry_uri):
    # Mock Unity Catalog behavior
    client = MlflowClient(registry_uri=registry_uri)

    # Mock the search_prompt_versions to return a PagedList with versions
    mock_versions = PagedList([Mock(version="1")], None)

    with (
        patch.object(client, "search_prompt_versions", return_value=mock_versions),
        patch.object(client, "_registry_uri", registry_uri),
    ):
        with pytest.raises(
            MlflowException, match=r"Cannot delete prompt .* because it still has undeleted"
        ):
            client.delete_prompt("test_prompt")


def test_link_prompt_version_to_model_smoke_test(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Create an experiment and a run to have a proper context
    experiment_id = client.create_experiment("test_experiment")
    with mlflow.start_run(experiment_id=experiment_id):
        # Create a model with a run context
        model = client.create_logged_model(experiment_id=experiment_id)

        # Register a prompt
        client.register_prompt(name="test_prompt", template="Hello, {{name}}!")

        # Link the prompt version to the model (this should not raise an exception)
        # This is the main assertion - that the method call succeeds
        client.link_prompt_version_to_model(
            name="test_prompt", version="1", model_id=model.model_id
        )


def test_link_prompts_to_trace_smoke_test(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Create an experiment and a run to have a proper context
    experiment_id = client.create_experiment("test_experiment")
    with mlflow.start_run(experiment_id=experiment_id):
        # Create a simple trace for testing
        trace_info = client.start_trace("test_trace")
        trace_id = trace_info.request_id

        # Register a prompt
        client.register_prompt(name="test_prompt", template="Hello, {{name}}!")

        # Get the prompt version and link to the trace (this should not raise an exception)
        # This is the main assertion - that the method call succeeds
        prompt_version = client.get_prompt_version("test_prompt", "1")
        client.link_prompt_versions_to_trace(prompt_versions=[prompt_version], trace_id=trace_id)


def test_log_model_artifact(tmp_path: Path, tracking_uri: str) -> None:
    client = MlflowClient(tracking_uri=tracking_uri)
    experiment_id = client.create_experiment("test")
    model = client.create_logged_model(experiment_id=experiment_id)
    tmp_path = tmp_path.joinpath("artifacts")
    tmp_path.mkdir()
    tmp_file = tmp_path.joinpath("file")
    tmp_file.write_text("a")
    client.log_model_artifact(model_id=model.model_id, local_path=str(tmp_file))
    artifacts = client.list_logged_model_artifacts(model_id=model.model_id)
    assert artifacts == [FileInfo(path="file", is_dir=False, file_size=1)]
    another_tmp_file = tmp_path.joinpath("another_file")
    another_tmp_file.write_text("aa")
    client.log_model_artifact(model_id=model.model_id, local_path=str(another_tmp_file))
    artifacts = client.list_logged_model_artifacts(model_id=model.model_id)
    artifacts = sorted(artifacts, key=lambda x: x.path)
    assert artifacts == [
        FileInfo(path="another_file", is_dir=False, file_size=2),
        FileInfo(path="file", is_dir=False, file_size=1),
    ]


def test_log_model_artifacts(tmp_path: Path, tracking_uri: str) -> None:
    client = MlflowClient(tracking_uri=tracking_uri)
    experiment_id = client.create_experiment("test")
    model = client.create_logged_model(experiment_id=experiment_id)
    tmp_path = tmp_path.joinpath("artifacts")
    tmp_path.mkdir()
    tmp_file = tmp_path.joinpath("file")
    tmp_file.write_text("a")
    tmp_dir = tmp_path.joinpath("dir")
    tmp_dir.mkdir()
    another_file = tmp_dir.joinpath("another_file")
    another_file.write_text("aa")
    client.log_model_artifacts(model_id=model.model_id, local_dir=str(tmp_path))
    artifacts = client.list_logged_model_artifacts(model_id=model.model_id)
    artifacts = sorted(artifacts, key=lambda x: x.path)
    assert artifacts == [
        FileInfo(path="dir", is_dir=True, file_size=None),
        FileInfo(path="file", is_dir=False, file_size=1),
    ]
    artifacts = client.list_logged_model_artifacts(model_id=model.model_id, path="dir")
    assert artifacts == [FileInfo(path="dir/another_file", is_dir=False, file_size=2)]


def test_logged_model_model_id_required(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    with pytest.raises(MlflowException, match="`model_id` must be a non-empty string, but got ''"):
        client.finalize_logged_model("", LoggedModelStatus.READY)

    with pytest.raises(MlflowException, match="`model_id` must be a non-empty string, but got ''"):
        client.get_logged_model("")

    with pytest.raises(MlflowException, match="`model_id` must be a non-empty string, but got ''"):
        client.delete_logged_model("")

    with pytest.raises(MlflowException, match="`model_id` must be a non-empty string, but got ''"):
        client.set_logged_model_tags("", {})

    with pytest.raises(MlflowException, match="`model_id` must be a non-empty string, but got ''"):
        client.delete_logged_model_tag("", "")


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
def test_log_metric_link_to_active_model(tracking_uri):
    model = mlflow.create_external_model(name="test_model")
    mlflow.set_active_model(name=model.name)
    client = MlflowClient(tracking_uri=tracking_uri)
    with mlflow.start_run() as run:
        client.log_metric(run.info.run_id, "metric", 1)
    logged_model = mlflow.get_logged_model(model_id=model.model_id)
    assert logged_model.name == model.name
    assert logged_model.model_id == model.model_id
    assert logged_model.metrics[0].key == "metric"
    assert logged_model.metrics[0].value == 1


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
def test_log_batch_link_to_active_model(tracking_uri):
    model = mlflow.create_external_model(name="test_model")
    mlflow.set_active_model(name=model.name)
    client = MlflowClient(tracking_uri=tracking_uri)
    with mlflow.start_run() as run:
        client.log_batch(run.info.run_id, [Metric("metric1", 1, 0, 0), Metric("metric2", 2, 0, 0)])
    logged_model = mlflow.get_logged_model(model_id=model.model_id)
    assert logged_model.name == model.name
    assert logged_model.model_id == model.model_id
    assert {m.key: m.value for m in logged_model.metrics} == {
        "metric1": 1,
        "metric2": 2,
    }


def test_load_prompt_with_alias_uri(tracking_uri, disable_prompt_cache):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Register two versions of a prompt
    client.register_prompt(name="alias_prompt", template="Hello, world!")
    client.register_prompt(name="alias_prompt", template="Hello, {{name}}!")

    # Assign alias to version 1
    client.set_prompt_alias("alias_prompt", alias="production", version=1)
    prompt = client.load_prompt("prompts:/alias_prompt@production")
    assert prompt.template == "Hello, world!"
    assert "production" in prompt.aliases

    # Reassign alias to version 2
    client.set_prompt_alias("alias_prompt", alias="production", version=2)
    prompt = client.load_prompt("prompts:/alias_prompt@production")
    assert prompt.template == "Hello, {{name}}!"
    assert "production" in prompt.aliases

    # Delete alias and verify loading fails
    client.delete_prompt_alias("alias_prompt", alias="production")
    with pytest.raises(
        MlflowException, match=r"Prompt (.*) does not exist.|Prompt alias (.*) not found."
    ):
        client.load_prompt("prompts:/alias_prompt@production")

    # Loading with the 'latest' alias
    prompt = client.load_prompt("prompts:/alias_prompt@latest")
    assert prompt.template == "Hello, {{name}}!"


def test_load_prompt_allow_missing_name_version(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Non-existent prompt by name+version should return None when allow_missing=True
    result = client.load_prompt("nonexistent_prompt", version=1, allow_missing=True)
    assert result is None

    # Non-existent prompt by name+version should raise exception when allow_missing=False
    with pytest.raises(MlflowException, match="Prompt with name=nonexistent_prompt not found"):
        client.load_prompt("nonexistent_prompt", version=1, allow_missing=False)

    # Existing prompt with non-existent version should return None when allow_missing=True
    client.register_prompt(name="existing_prompt", template="Hello, world!")
    result = client.load_prompt("existing_prompt", version=999, allow_missing=True)
    assert result is None

    # Existing prompt with non-existent version should raise exception when allow_missing=False
    with pytest.raises(
        MlflowException, match=r"Prompt \(name=existing_prompt, version=999\) not found"
    ):
        client.load_prompt("existing_prompt", version=999, allow_missing=False)


def test_load_prompt_allow_missing_uri_version(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Non-existent prompt by URI+version should return None when allow_missing=True
    result = client.load_prompt("prompts:/nonexistent_prompt/1", allow_missing=True)
    assert result is None

    # Non-existent prompt by URI+version should raise exception when allow_missing=False
    with pytest.raises(MlflowException, match="Prompt with name=nonexistent_prompt not found"):
        client.load_prompt("prompts:/nonexistent_prompt/1", allow_missing=False)

    # Existing prompt with non-existent version via URI should return None when allow_missing=True
    client.register_prompt(name="existing_prompt", template="Hello, world!")
    result = client.load_prompt("prompts:/existing_prompt/999", allow_missing=True)
    assert result is None

    # Existing prompt with non-existent version via URI should raise when allow_missing=False
    with pytest.raises(
        MlflowException, match=r"Prompt \(name=existing_prompt, version=999\) not found"
    ):
        client.load_prompt("prompts:/existing_prompt/999", allow_missing=False)


def test_load_prompt_allow_missing_uri_alias(tracking_uri):
    client = MlflowClient(tracking_uri=tracking_uri)

    # Non-existent prompt with alias should return None when allow_missing=True
    result = client.load_prompt("prompts:/nonexistent_prompt@production", allow_missing=True)
    assert result is None

    # Non-existent prompt with alias should raise exception when allow_missing=False
    with pytest.raises(MlflowException, match="Prompt with name=nonexistent_prompt not found"):
        client.load_prompt("prompts:/nonexistent_prompt@production", allow_missing=False)

    # Existing prompt with non-existent alias should return None when allow_missing=True
    client.register_prompt(name="existing_prompt", template="Hello, world!")
    result = client.load_prompt("prompts:/existing_prompt@nonexistent_alias", allow_missing=True)
    assert result is None

    # Existing prompt with non-existent alias should raise exception when allow_missing=False
    with pytest.raises(MlflowException, match="Prompt alias nonexistent_alias not found"):
        client.load_prompt("prompts:/existing_prompt@nonexistent_alias", allow_missing=False)


def test_create_prompt_chat_format_client_integration():
    chat_template = [
        {"role": "system", "content": "You are a {{style}} assistant."},
        {"role": "user", "content": "{{question}}"},
    ]

    response_format = {"type": "string"}

    # Use client to create prompt
    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_chat_client",
        template=chat_template,
        response_format=response_format,
        commit_message="Test chat prompt via client",
    )

    assert prompt.template == chat_template
    assert prompt.response_format == response_format

    # Load via client
    loaded_prompt = client.get_prompt_version("test_chat_client", 1)
    assert not loaded_prompt.is_text_prompt
    assert loaded_prompt.template == chat_template
    assert loaded_prompt.response_format == response_format


def test_link_chat_prompt_version_to_run():
    chat_template = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello {{name}}!"},
    ]

    client = MlflowClient()
    prompt = client.register_prompt(name="test_chat_link", template=chat_template)

    # Create run and link prompt
    run = client.create_run(client.create_experiment("test_exp"))
    client.link_prompt_version_to_run(run.info.run_id, prompt)

    # Verify linking
    run_data = client.get_run(run.info.run_id)
    linked_prompts_tag = run_data.data.tags.get(TraceTagKey.LINKED_PROMPTS)
    assert linked_prompts_tag is not None

    linked_prompts = json.loads(linked_prompts_tag)
    assert len(linked_prompts) == 1
    assert linked_prompts[0]["name"] == "test_chat_link"
    assert linked_prompts[0]["version"] == "1"


def test_create_prompt_with_pydantic_response_format_client():
    class ResponseSchema(BaseModel):
        answer: str
        confidence: float

    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_pydantic_client",
        template="What is {{question}}?",
        response_format=ResponseSchema,
        commit_message="Test Pydantic response format via client",
    )

    assert prompt.response_format == ResponseSchema.model_json_schema()
    assert prompt.commit_message == "Test Pydantic response format via client"

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_pydantic_client", 1)
    assert loaded_prompt.response_format == ResponseSchema.model_json_schema()


def test_create_prompt_with_dict_response_format_client():
    response_format = {
        "type": "object",
        "properties": {
            "summary": {"type": "string"},
            "key_points": {"type": "array", "items": {"type": "string"}},
        },
    }

    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_dict_response_client",
        template="Analyze this: {{text}}",
        response_format=response_format,
        tags={"analysis_type": "text"},
    )

    assert prompt.response_format == response_format
    assert prompt.tags["analysis_type"] == "text"

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_dict_response_client", 1)
    assert loaded_prompt.response_format == response_format


def test_create_prompt_text_backward_compatibility_client():
    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_text_backward_client",
        template="Hello {{name}}!",
        commit_message="Test backward compatibility via client",
    )

    assert prompt.is_text_prompt
    assert prompt.template == "Hello {{name}}!"
    assert prompt.commit_message == "Test backward compatibility via client"

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_text_backward_client", 1)
    assert loaded_prompt.is_text_prompt
    assert loaded_prompt.template == "Hello {{name}}!"


def test_create_prompt_complex_chat_template_client():
    chat_template = [
        {
            "role": "system",
            "content": "You are a {{style}} assistant named {{name}}.",
        },
        {"role": "user", "content": "{{greeting}}! {{question}}"},
        {
            "role": "assistant",
            "content": "I understand you're asking about {{topic}}.",
        },
    ]

    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_complex_chat_client",
        template=chat_template,
        tags={"complexity": "high"},
    )

    assert prompt.template == chat_template
    assert prompt.tags["complexity"] == "high"

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_complex_chat_client", 1)
    assert not loaded_prompt.is_text_prompt
    assert loaded_prompt.template == chat_template


def test_create_prompt_with_none_response_format_client():
    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_none_response_client",
        template="Hello {{name}}!",
        response_format=None,
    )

    assert prompt.response_format is None

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_none_response_client", 1)
    assert loaded_prompt.response_format is None


def test_create_prompt_with_single_message_chat_client():
    chat_template = [{"role": "user", "content": "Hello {{name}}!"}]

    client = MlflowClient()
    prompt = client.register_prompt(name="test_single_message_client", template=chat_template)

    assert prompt.template == chat_template
    assert prompt.variables == {"name"}

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_single_message_client", 1)
    assert not loaded_prompt.is_text_prompt
    assert loaded_prompt.template == chat_template


def test_create_prompt_with_multiple_variables_in_chat_client():
    chat_template = [
        {
            "role": "system",
            "content": "You are a {{style}} assistant named {{name}}.",
        },
        {"role": "user", "content": "{{greeting}}! {{question}}"},
        {
            "role": "assistant",
            "content": "I understand you're asking about {{topic}}.",
        },
    ]

    client = MlflowClient()
    prompt = client.register_prompt(name="test_multiple_variables_client", template=chat_template)

    expected_variables = {"style", "name", "greeting", "question", "topic"}
    assert prompt.variables == expected_variables

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_multiple_variables_client", 1)
    assert loaded_prompt.variables == expected_variables


def test_create_prompt_with_mixed_content_types_client():
    chat_template = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello {{name}}!"},
        {"role": "assistant", "content": "Hi there! How can I help you today?"},
    ]

    client = MlflowClient()
    prompt = client.register_prompt(name="test_mixed_content_client", template=chat_template)

    assert prompt.template == chat_template
    assert prompt.variables == {"name"}

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_mixed_content_client", 1)
    assert not loaded_prompt.is_text_prompt
    assert loaded_prompt.template == chat_template


def test_create_prompt_with_nested_variables_client():
    chat_template = [
        {
            "role": "system",
            "content": "You are a {{user.preferences.style}} assistant.",
        },
        {
            "role": "user",
            "content": "Hello {{user.name}}! {{user.preferences.greeting}}",
        },
    ]

    client = MlflowClient()
    prompt = client.register_prompt(name="test_nested_variables_client", template=chat_template)

    expected_variables = {
        "user.preferences.style",
        "user.name",
        "user.preferences.greeting",
    }
    assert prompt.variables == expected_variables

    # Load and verify
    loaded_prompt = client.get_prompt_version("test_nested_variables_client", 1)
    assert loaded_prompt.variables == expected_variables


def test_link_prompt_with_response_format_to_run():
    response_format = {
        "type": "object",
        "properties": {"answer": {"type": "string"}},
    }
    client = MlflowClient()
    prompt = client.register_prompt(
        name="test_response_link",
        template="What is {{question}}?",
        response_format=response_format,
    )

    # Create run and link prompt
    run = client.create_run(client.create_experiment("test_exp"))
    client.link_prompt_version_to_run(run.info.run_id, prompt)

    # Verify linking
    run_data = client.get_run(run.info.run_id)
    linked_prompts_tag = run_data.data.tags.get(TraceTagKey.LINKED_PROMPTS)
    assert linked_prompts_tag is not None

    linked_prompts = json.loads(linked_prompts_tag)
    assert len(linked_prompts) == 1
    assert linked_prompts[0]["name"] == "test_response_link"
    assert linked_prompts[0]["version"] == "1"


def test_link_multiple_prompt_types_to_run():
    client = MlflowClient()

    # Create text prompt
    text_prompt = client.register_prompt(name="test_text_link", template="Hello {{name}}!")

    # Create chat prompt
    chat_template = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "{{question}}"},
    ]
    chat_prompt = client.register_prompt(name="test_chat_link_multiple", template=chat_template)

    # Create run and link both prompts
    run = client.create_run(client.create_experiment("test_exp"))
    client.link_prompt_version_to_run(run.info.run_id, text_prompt)
    client.link_prompt_version_to_run(run.info.run_id, chat_prompt)

    # Verify linking
    run_data = client.get_run(run.info.run_id)
    linked_prompts_tag = run_data.data.tags.get(TraceTagKey.LINKED_PROMPTS)
    assert linked_prompts_tag is not None

    linked_prompts = json.loads(linked_prompts_tag)
    assert len(linked_prompts) == 2

    expected_prompts = [
        {"name": "test_text_link", "version": "1"},
        {"name": "test_chat_link_multiple", "version": "1"},
    ]
    for expected_prompt in expected_prompts:
        assert expected_prompt in linked_prompts


def test_mlflow_client_create_dataset(mock_store):
    created_dataset = EvaluationDataset(
        dataset_id="test_dataset_id",
        name="test_dataset",
        digest="abcdef123456",
        created_time=1234567890,
        last_update_time=1234567890,
        tags={"environment": "production", "version": "1.0"},
    )
    created_dataset.experiment_ids = ["exp1", "exp2"]
    mock_store.create_dataset.return_value = created_dataset

    # Mock context registry to return empty tags so mlflow.user is not auto-added
    with mock.patch(
        "mlflow.tracking._tracking_service.client.context_registry.resolve_tags", return_value={}
    ):
        dataset = MlflowClient().create_dataset(
            name="qa_evaluation",
            experiment_id=["exp1", "exp2"],
            tags={"environment": "production", "version": "1.0"},
        )

    assert dataset.dataset_id == "test_dataset_id"
    assert dataset.name == "test_dataset"
    assert dataset.tags == {"environment": "production", "version": "1.0"}

    mock_store.create_dataset.assert_called_once_with(
        name="qa_evaluation",
        tags={"environment": "production", "version": "1.0"},
        experiment_ids=["exp1", "exp2"],
    )


def test_mlflow_client_create_evaluation_dataset_minimal(mock_store):
    created_dataset = EvaluationDataset(
        dataset_id="test_dataset_id",
        name="test_dataset",
        digest="abcdef123456",
        created_time=1234567890,
        last_update_time=1234567890,
    )
    mock_store.create_dataset.return_value = created_dataset

    # Mock context registry to return empty tags so mlflow.user is not auto-added
    with mock.patch(
        "mlflow.tracking._tracking_service.client.context_registry.resolve_tags", return_value={}
    ):
        dataset = MlflowClient().create_dataset(name="test_dataset")

    assert dataset.dataset_id == "test_dataset_id"
    assert dataset.name == "test_dataset"

    mock_store.create_dataset.assert_called_once_with(
        name="test_dataset",
        tags=None,
        experiment_ids=None,
    )


def test_mlflow_client_get_dataset(mock_store):
    mock_store.get_dataset.return_value = EvaluationDataset(
        dataset_id="dataset_123",
        name="test_dataset",
        digest="abcdef123456",
        created_time=1234567890,
        last_update_time=1234567890,
        tags={"source": "human-annotated"},
    )

    dataset = MlflowClient().get_dataset("dataset_123")

    assert dataset.dataset_id == "dataset_123"
    assert dataset.name == "test_dataset"
    assert dataset.tags == {"source": "human-annotated"}

    mock_store.get_dataset.assert_called_once_with("dataset_123")


def test_mlflow_client_delete_dataset(mock_store):
    MlflowClient().delete_dataset("dataset_123")

    mock_store.delete_dataset.assert_called_once_with("dataset_123")


def test_mlflow_client_search_datasets(mock_store):
    mock_store.search_datasets.return_value = PagedList(
        [
            EvaluationDataset(
                dataset_id="dataset_1",
                name="dataset_1",
                digest="digest1",
                created_time=1234567890,
                last_update_time=1234567890,
            ),
            EvaluationDataset(
                dataset_id="dataset_2",
                name="dataset_2",
                digest="digest2",
                created_time=1234567890,
                last_update_time=1234567890,
            ),
        ],
        "next_token",
    )

    result = MlflowClient().search_datasets(
        experiment_ids=["exp1", "exp2"],
        filter_string="name LIKE 'qa_%'",
        max_results=100,
        order_by=["created_time DESC"],
        page_token="page_token_123",
    )

    assert len(result) == 2
    assert result[0].dataset_id == "dataset_1"
    assert result[1].dataset_id == "dataset_2"
    assert result.token == "next_token"

    mock_store.search_datasets.assert_called_once_with(
        experiment_ids=["exp1", "exp2"],
        filter_string="name LIKE 'qa_%'",
        max_results=100,
        order_by=["created_time DESC"],
        page_token="page_token_123",
    )


def test_mlflow_client_search_datasets_empty_results(mock_store):
    mock_store.search_datasets.return_value = PagedList([], None)

    result = MlflowClient().search_datasets(
        experiment_ids=["exp1"], filter_string="name = 'nonexistent'"
    )

    assert len(result) == 0
    assert result.token is None


def test_mlflow_client_search_datasets_defaults(mock_store):
    mock_store.search_datasets.return_value = PagedList([], None)

    result = MlflowClient().search_datasets()

    assert len(result) == 0
    assert result.token is None

    mock_store.search_datasets.assert_called_once_with(
        experiment_ids=None,
        filter_string=None,
        max_results=SEARCH_EVALUATION_DATASETS_MAX_RESULTS,
        order_by=None,
        page_token=None,
    )


@pytest.mark.skipif(is_windows(), reason="FileStore URI handling issues on Windows")
def test_mlflow_client_datasets_filestore_not_supported(tmp_path):
    pytest.skip("FileStore is no longer supported.")
    file_store_uri = str(tmp_path)
    client = MlflowClient(tracking_uri=file_store_uri)

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.create_dataset(name="test_dataset")
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.get_dataset("dataset_123")
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.delete_dataset("dataset_123")
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.search_datasets()
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.set_dataset_tags("dataset_123", {"tag1": "value1"})
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.delete_dataset_tag("dataset_123", "tag1")
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.add_dataset_to_experiments("dataset_123", ["1", "2"])
    assert exc_info.value.error_code == "FEATURE_DISABLED"

    with pytest.raises(MlflowException, match="is not supported with FileStore") as exc_info:
        client.remove_dataset_from_experiments("dataset_123", ["1", "2"])
    assert exc_info.value.error_code == "FEATURE_DISABLED"


def test_mlflow_client_set_dataset_tags(mock_store):
    MlflowClient().set_dataset_tags(
        dataset_id="dataset_123",
        tags={"env": "prod", "version": "2.0"},
    )

    mock_store.set_dataset_tags.assert_called_once_with(
        dataset_id="dataset_123",
        tags={"env": "prod", "version": "2.0"},
    )


def test_mlflow_client_delete_dataset_tag(mock_store):
    MlflowClient().delete_dataset_tag(
        dataset_id="dataset_123",
        key="deprecated",
    )

    mock_store.delete_dataset_tag.assert_called_once_with(
        dataset_id="dataset_123",
        key="deprecated",
    )


def test_mlflow_client_add_dataset_to_experiments(mock_store):
    mock_dataset = Mock(spec=EvaluationDataset)
    mock_dataset.dataset_id = "dataset_123"
    mock_dataset.experiment_ids = ["1", "2", "3"]
    mock_store.add_dataset_to_experiments.return_value = mock_dataset

    client = MlflowClient()
    result = client.add_dataset_to_experiments(
        dataset_id="dataset_123",
        experiment_ids=["2", "3"],
    )

    assert result == mock_dataset
    assert result.experiment_ids == ["1", "2", "3"]
    mock_store.add_dataset_to_experiments.assert_called_once_with("dataset_123", ["2", "3"])


def test_mlflow_client_remove_dataset_from_experiments(mock_store):
    mock_dataset = Mock(spec=EvaluationDataset)
    mock_dataset.dataset_id = "dataset_123"
    mock_dataset.experiment_ids = ["1"]
    mock_store.remove_dataset_from_experiments.return_value = mock_dataset

    client = MlflowClient()
    result = client.remove_dataset_from_experiments(
        dataset_id="dataset_123",
        experiment_ids=["2", "3"],
    )

    assert result == mock_dataset
    assert result.experiment_ids == ["1"]
    mock_store.remove_dataset_from_experiments.assert_called_once_with("dataset_123", ["2", "3"])


def test_mlflow_client_dataset_associations_databricks_blocking(mock_store):
    with mock.patch("mlflow.utils.databricks_utils.is_databricks_uri") as mock_is_dbx:
        mock_is_dbx.return_value = True
        client = MlflowClient(tracking_uri="databricks")

        with pytest.raises(
            MlflowException, match="not supported when tracking URI is 'databricks'"
        ) as exc_info:
            client.add_dataset_to_experiments("dataset_123", ["1", "2"])
        assert exc_info.value.error_code == "INVALID_PARAMETER_VALUE"

        with pytest.raises(
            MlflowException, match="not supported when tracking URI is 'databricks'"
        ) as exc_info:
            client.remove_dataset_from_experiments("dataset_123", ["1", "2"])
        assert exc_info.value.error_code == "INVALID_PARAMETER_VALUE"


def test_log_spans_and_get_trace_with_sqlalchemy_store(tmp_path: Path) -> None:
    tracking_uri = f"sqlite:///{tmp_path}/test.db"

    with _use_tracking_uri(tracking_uri):
        client = MlflowClient()

        assert isinstance(client._tracking_client.store, SqlAlchemyTrackingStore)

        experiment_id = client.create_experiment("test_log_spans_get_trace")
        trace_id = f"tr-{uuid.uuid4().hex}"

        # Create test spans using OpenTelemetry format
        otel_span1 = OTelReadableSpan(
            name="parent_span",
            context=trace_api.SpanContext(
                trace_id=12345,
                span_id=111,
                is_remote=False,
                trace_flags=trace_api.TraceFlags(1),
            ),
            parent=None,
            attributes={
                "mlflow.traceRequestId": json.dumps(trace_id, cls=TraceJSONEncoder),
                "llm.model_name": "test-model",
                "custom.attribute": "parent-value",
            },
            start_time=1_000_000_000,
            end_time=2_000_000_000,
            resource=None,
        )

        otel_span2 = OTelReadableSpan(
            name="child_span",
            context=trace_api.SpanContext(
                trace_id=12345,
                span_id=222,
                is_remote=False,
                trace_flags=trace_api.TraceFlags(1),
            ),
            parent=trace_api.SpanContext(
                trace_id=12345,
                span_id=111,
                is_remote=False,
                trace_flags=trace_api.TraceFlags(1),
            ),
            attributes={
                "mlflow.traceRequestId": json.dumps(trace_id, cls=TraceJSONEncoder),
                "operation.type": "database_query",
                "custom.attribute": "child-value",
            },
            start_time=1_200_000_000,
            end_time=1_800_000_000,
            resource=None,
        )

        # Convert to MLflow spans
        mlflow_spans = [
            create_mlflow_span(otel_span1, trace_id, "LLM"),
            create_mlflow_span(otel_span2, trace_id, "LLM"),
        ]

        # Log spans directly to the store (simulating OTLP endpoint)
        store = client._tracking_client.store
        logged_spans = store.log_spans(experiment_id, mlflow_spans)

        # Verify spans were logged
        assert len(logged_spans) == 2

        # Verify the trace has the spans location tag set
        trace_info = store.get_trace_info(trace_id)
        assert trace_info.tags.get(TraceTagKey.SPANS_LOCATION) == SpansLocation.TRACKING_STORE

        # Now test that mlflow.get_trace() works and loads spans from the database
        trace = mlflow.get_trace(trace_id)

        # Verify trace structure
        assert trace.info.trace_id == trace_id
        assert trace.info.tags.get(TraceTagKey.SPANS_LOCATION) == SpansLocation.TRACKING_STORE

        # Verify spans were loaded from database
        assert len(trace.data.spans) == 2

        # Sort spans by start time for consistent testing
        spans_by_start_time = sorted(trace.data.spans, key=lambda s: s.start_time_ns)

        # Verify parent span
        parent_span = spans_by_start_time[0]
        assert parent_span.name == "parent_span"
        assert parent_span.trace_id == trace_id
        assert parent_span.start_time_ns == 1_000_000_000
        assert parent_span.end_time_ns == 2_000_000_000
        assert parent_span.attributes.get("llm.model_name") == "test-model"
        assert parent_span.attributes.get("custom.attribute") == "parent-value"

        # Verify child span
        child_span = spans_by_start_time[1]
        assert child_span.name == "child_span"
        assert child_span.trace_id == trace_id
        assert child_span.start_time_ns == 1_200_000_000
        assert child_span.end_time_ns == 1_800_000_000
        assert child_span.attributes.get("operation.type") == "database_query"
        assert child_span.attributes.get("custom.attribute") == "child-value"


def test_mlflow_get_trace_with_sqlalchemy_store(tmp_path: Path) -> None:
    tracking_uri = f"sqlite:///{tmp_path}/test.db"

    with _use_tracking_uri(tracking_uri):
        client = MlflowClient()

        assert isinstance(client._tracking_client.store, SqlAlchemyTrackingStore)

        with mlflow.start_span() as span:
            pass

        trace_id = span.trace_id
        mlflow.flush_trace_async_logging()
        sql_alchemy_store_module = "mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore"
        with (
            mock.patch(f"{sql_alchemy_store_module}.get_trace") as mock_get_trace,
        ):
            mlflow.get_trace(trace_id)

        mock_get_trace.assert_called_once_with(trace_id)

        with (
            mock.patch(
                f"{sql_alchemy_store_module}.get_trace",
                side_effect=MlflowNotImplementedException,
            ) as mock_get_trace,
            mock.patch(f"{sql_alchemy_store_module}.batch_get_traces") as mock_batch_get_traces,
        ):
            mlflow.get_trace(trace_id)

        mock_get_trace.assert_called_once_with(trace_id)
        mock_batch_get_traces.assert_called_once_with([trace_id])


def test_create_issue_basic(tmp_path: Path):
    tracking_uri = f"sqlite:///{tmp_path}/test.db"

    with _use_tracking_uri(tracking_uri):
        client = MlflowClient()
        exp_id = client.create_experiment("test_create_issue")
        tracing_client = client._tracing_client

        issue = tracing_client._create_issue(
            experiment_id=exp_id,
            name="Test issue",
            description="This is a test issue",
        )

        assert issue.issue_id.startswith("iss-")
        assert issue.experiment_id == exp_id
        assert issue.name == "Test issue"
        assert issue.description == "This is a test issue"
        assert issue.status == IssueStatus.PENDING
        assert issue.severity is None
        assert issue.root_causes is None
        assert issue.source_run_id is None
        assert issue.created_by is None
        assert issue.created_timestamp > 0
        assert issue.last_updated_timestamp == issue.created_timestamp


def test_create_issue_with_all_fields(tmp_path: Path):
    tracking_uri = f"sqlite:///{tmp_path}/test.db"

    with _use_tracking_uri(tracking_uri):
        client = MlflowClient()
        exp_id = client.create_experiment("test_create_issue_all_fields")
        tracing_client = client._tracing_client
        with mlflow.start_run(experiment_id=exp_id) as run:
            issue = tracing_client._create_issue(
                experiment_id=exp_id,
                name="High latency",
                description="API response times exceed threshold",
                status=IssueStatus.RESOLVED,
                severity=IssueSeverity.HIGH,
                root_causes=["Database query slow", "Network congestion"],
                source_run_id=run.info.run_id,
                created_by="monitoring_system",
            )

    assert issue.issue_id.startswith("iss-")
    assert issue.experiment_id == exp_id
    assert issue.name == "High latency"
    assert issue.description == "API response times exceed threshold"
    assert issue.status == IssueStatus.RESOLVED
    assert issue.severity == IssueSeverity.HIGH
    assert issue.root_causes == ["Database query slow", "Network congestion"]
    assert issue.source_run_id == run.info.run_id
    assert issue.created_by == "monitoring_system"
    assert issue.created_timestamp > 0