1  from pathlib import Path
 2  
 3  from mlflow.environment_variables import MLFLOW_TRACKING_PASSWORD, MLFLOW_TRACKING_USERNAME
 4  from mlflow.server.auth import auth_config
 5  
 6  from tests.helper_functions import random_str
 7  from tests.tracking.integration_test_utils import _send_rest_tracking_post_request
 8  
 9  PERMISSION = "READ"
10  NEW_PERMISSION = "EDIT"
11  ADMIN_USERNAME = auth_config.admin_username
12  ADMIN_PASSWORD = auth_config.admin_password
13  
14  
15  def write_isolated_auth_config(tmp_path: Path) -> Path:
16      """Write a basic_auth.ini under ``tmp_path`` whose ``database_uri`` points
17      at an SQLite file inside the same dir, and return the config path.
18  
19      Tests that spawn the auth server via ``_init_server`` must point
20      ``MLFLOW_AUTH_CONFIG_PATH`` at the returned file (through ``extra_env``) so
21      the spawned auth server writes to the temp DB instead of the repo-root
22      ``basic_auth.db`` shared with the dev server. Without this, integration
23      tests pollute (and depending on the fixture, delete) the developer's local
24      auth state.
25      """
26      config_path = tmp_path / "basic_auth.ini"
27      db_path = tmp_path / "basic_auth.db"
28      config_path.write_text(
29          "[mlflow]\n"
30          "default_permission = READ\n"
31          f"database_uri = sqlite:///{db_path}\n"
32          f"admin_username = {ADMIN_USERNAME}\n"
33          f"admin_password = {ADMIN_PASSWORD}\n"
34          "authorization_function = mlflow.server.auth:authenticate_request_basic_auth\n"
35          "grant_default_workspace_access = false\n"
36      )
37      return config_path
38  
39  
40  def create_user(tracking_uri: str, username: str | None = None, password: str | None = None):
41      username = random_str() if username is None else username
42      password = random_str() if password is None else password
43      response = _send_rest_tracking_post_request(
44          tracking_uri,
45          "/api/2.0/mlflow/users/create",
46          {
47              "username": username,
48              "password": password,
49          },
50          auth=(ADMIN_USERNAME, ADMIN_PASSWORD),
51      )
52      response.raise_for_status()
53      return username, password
54  
55  
56  class User:
57      def __init__(self, username, password, monkeypatch):
58          self.username = username
59          self.password = password
60          self.monkeypatch = monkeypatch
61  
62      def __enter__(self):
63          self.monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, self.username)
64          self.monkeypatch.setenv(MLFLOW_TRACKING_PASSWORD.name, self.password)
65  
66      def __exit__(self, exc_type, exc_val, exc_tb):
67          self.monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False)
68          self.monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False)