Add basic auth integration test

docker-compose.aws-mysql-test.yaml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   minio:
     container_name: minio-mysql-test

docker-compose.aws-postgres-test.yaml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   minio:
     container_name: minio-pg-test

docker-compose.azure-mysql-test.yaml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   azurite:
     container_name: azurite-mysql-test

docker-compose.azure-postgres-test.yaml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   azurite:
     container_name: azurite-pg-test

docker-compose.basic-auth-test.yaml

@@ -0,0 +1,87 @@
+services:
+  minio:
+    container_name: minio-basic-auth-test
+    hostname: minio
+    image: minio/minio:${MINIO_VERSION}
+    entrypoint: sh
+    command: -c 'mkdir -p /data/mlflow && minio server /data --console-address ":9001"'
+    environment:
+      MINIO_ROOT_USER: minioadmin
+      MINIO_ROOT_PASSWORD: minioadmin
+    ports:
+      - "9000:9000"
+      - "9001:9001"
+    volumes:
+      - minio-basic-auth-storage:/data
+    networks:
+      - basic_auth_test_nw
+
+  postgres:
+    image: "postgres:${POSTGRES_VERSION}"
+    container_name: mlflow-basic-auth-db
+    environment:
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_USER: postgres
+      POSTGRES_DB: mlflow
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres-basic-auth-storage:/var/lib/postgresql/data
+    networks:
+      - basic_auth_test_nw
+
+  usersdb:
+    image: "postgres:${POSTGRES_VERSION}"
+    container_name: mlflow-basic-auth-users-db
+    environment:
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_USER: postgres
+      POSTGRES_DB: users
+    ports:
+      - "5433:5432"
+    volumes:
+      - postgres-basic-auth-users-storage:/var/lib/postgresql/data
+    networks:
+      - basic_auth_test_nw
+
+  mlflow:
+    container_name: mlflow-basic-auth-test
+    build:
+      context: .
+      dockerfile: Dockerfile
+    command: "mlflow server --backend-store-uri=postgresql:// --default-artifact-root=s3://mlflow/ --host=0.0.0.0 --port=8080 --app-name basic-auth"
+    environment:
+      MLFLOW_S3_ENDPOINT_URL: http://minio:9000
+      AWS_ACCESS_KEY_ID: minioadmin
+      AWS_SECRET_ACCESS_KEY: minioadmin
+      PGHOST: postgres
+      PGPORT: 5432
+      PGDATABASE: mlflow
+      PGUSER: postgres
+      PGPASSWORD: postgres
+      MLFLOW_AUTH_CONFIG_PATH: /mlflow/basic_auth.ini
+    ports:
+      - "8080:8080"
+    networks:
+      - basic_auth_test_nw
+    volumes:
+      - mlflow-basic-auth-storage:/mlflow
+      - ./test-containers/basic-auth/basic_auth.ini:/mlflow/basic_auth.ini
+    depends_on:
+      - minio
+      - postgres
+      - usersdb
+
+volumes:
+  postgres-basic-auth-storage:
+    driver: local
+  postgres-basic-auth-users-storage:
+    driver: local
+  mlflow-basic-auth-storage:
+    driver: local
+  minio-basic-auth-storage:
+    driver: local
+
+networks:
+  basic_auth_test_nw:
+    driver: bridge

docker-compose.gcp-mysql-test.yaml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   gcs:
     container_name: gcs-mysql-test

docker-compose.gcp-postgres-test.yaml

@@ -1,5 +1,3 @@
-version: "3.8"
-
 services:
   gcs:
     container_name: gcs-pg-test

test-containers/basic-auth/basic_auth.ini

@@ -0,0 +1,7 @@
+[mlflow]
+default_permission = READ # Default permission for all users. More details: https://mlflow.org/docs/latest/auth/index.html#permissions
+# database_uri = sqlite:///basic_auth.db
+database_uri = postgresql://postgres:postgres@usersdb:5432/users
+admin_username = testuser
+admin_password = simpletestpassword
+authorization_function = mlflow.server.auth:authenticate_request_basic_auth

tests/extended_docker_compose.py

@@ -0,0 +1,39 @@
+import re
+import subprocess
+import time
+
+from testcontainers.compose import DockerCompose
+
+
+class ExtendedDockerCompose(DockerCompose):
+    def wait_for_logs(self, service_name, expected_log, timeout=120, interval=5):
+        # Get the directory where the compose file is located (context)
+        compose_dir = self.context
+        compose_files = self.compose_file_name
+
+        if isinstance(compose_files, list):
+            compose_file_args = []
+            for file in compose_files:
+                compose_file_args.extend(["-f", file])
+        else:
+            compose_file_args = ["-f", compose_files]
+
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            try:
+                # Fetch the logs of the specific service using subprocess
+                logs = subprocess.check_output(
+                    ["docker-compose", *compose_file_args, "logs", service_name],
+                    cwd=compose_dir,  # Use the correct directory
+                    text=True,
+                )
+                # Use regex search to match the expected log pattern
+                if re.search(expected_log, logs):
+                    print(f"Found expected log matching: {expected_log}")
+                    return
+            except subprocess.CalledProcessError as e:
+                print(f"Error fetching logs: {e}")
+            time.sleep(interval)
+        raise TimeoutError(
+            f"Log message matching '{expected_log!r}' not found within {timeout} seconds."
+        )

tests/test_basic_auth_postgres_integration.py

@@ -0,0 +1,74 @@
+import os
+import re
+
+import requests
+
+from .extended_docker_compose import ExtendedDockerCompose
+
+import mlflow
+from mlflow.tracking.client import MlflowClient
+
+# from testcontainers.compose import DockerCompose
+# from testcontainers.core.waiting_utils import wait_for_logs
+
+
+def test_postgres_backended_model_upload_and_access_with_basic_auth(
+    test_model, training_params, conda_env
+):
+    with ExtendedDockerCompose(
+        context=".",
+        compose_file_name=["docker-compose.basic-auth-test.yaml"],
+        # pull=True,
+        build=True,
+    ) as compose:
+        mlflow_host = compose.get_service_host("mlflow", 8080)
+        mlflow_port = compose.get_service_port("mlflow", 8080)
+        minio_host = compose.get_service_host("minio", 9000)
+        minio_port = compose.get_service_port("minio", 9000)
+
+        admin_username = "testuser"
+        admin_password = "simpletestpassword"
+
+        base_url = f"http://{mlflow_host}:{mlflow_port}"
+
+        log_message = f".*Listening at: {re.escape(base_url)}.*"
+        compose.wait_for_logs("mlflow", log_message)
+
+        experiment_name = "aws-cloud-postgres-experiment"
+        model_name = "test-aws-pg-model"
+        stage_name = "Staging"
+        os.environ["MLFLOW_TRACKING_USERNAME"] = admin_username
+        os.environ["MLFLOW_TRACKING_PASSWORD"] = admin_password
+
+        mlflow.set_tracking_uri(base_url)
+        mlflow.set_experiment(experiment_name)
+        experiment = mlflow.get_experiment_by_name(experiment_name)
+
+        os.environ["MLFLOW_S3_ENDPOINT_URL"] = f"http://{minio_host}:{minio_port}"
+        os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
+        os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"
+
+        with mlflow.start_run(experiment_id=experiment.experiment_id) as run:
+            mlflow.log_params(training_params)
+            mlflow.pyfunc.log_model("model", conda_env=conda_env, python_model=test_model)
+            model_uri = f"runs:/{run.info.run_id}/model"
+            model_details = mlflow.register_model(model_uri, model_name)
+
+            client = MlflowClient()
+            client.set_registered_model_alias(
+                name=model_details.name,
+                alias=stage_name,
+                version=model_details.version,
+            )
+
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(
+            url=latest_version_url,
+            params=params,
+            auth=(admin_username, admin_password),
+            timeout=300,
+        )
+
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]