1  from __future__ import annotations
 2  
 3  import os
 4  
 5  from mlflow.environment_variables import MLFLOW_WORKSPACE
 6  from mlflow.utils.workspace_context import WorkspaceContext, clear_server_request_workspace
 7  from mlflow.utils.workspace_utils import (
 8      DEFAULT_WORKSPACE_NAME,
 9      resolve_entity_workspace_name,
10  )
11  
12  
13  def teardown_function():
14      # Ensure the ContextVar does not leak between tests
15      clear_server_request_workspace()
16      os.environ.pop(MLFLOW_WORKSPACE.name, None)
17  
18  
19  def test_resolve_entity_workspace_prefers_argument():
20      assert resolve_entity_workspace_name("  team-arg  ") == "team-arg"
21  
22  
23  def test_resolve_entity_workspace_uses_context_var():
24      with WorkspaceContext("ctx-workspace"):
25          assert resolve_entity_workspace_name(None) == "ctx-workspace"
26  
27  
28  def test_resolve_entity_workspace_falls_back_to_env(monkeypatch):
29      monkeypatch.delenv(MLFLOW_WORKSPACE.name, raising=False)
30      with WorkspaceContext(None):
31          pass
32      monkeypatch.setenv(MLFLOW_WORKSPACE.name, "  env-workspace  ")
33      assert resolve_entity_workspace_name(None) == "env-workspace"
34  
35  
36  def test_resolve_entity_workspace_defaults_when_unset(monkeypatch):
37      monkeypatch.delenv(MLFLOW_WORKSPACE.name, raising=False)
38      with WorkspaceContext(None):
39          pass
40      assert resolve_entity_workspace_name(None) == DEFAULT_WORKSPACE_NAME