From 6317c6c57b95bfec58914241631e696a0c5d2a33 Mon Sep 17 00:00:00 2001 From: Burak Ince Date: Sun, 15 Dec 2024 10:45:32 +0100 Subject: [PATCH] Add basic auth integration test --- docker-compose.aws-mysql-test.yaml | 2 - docker-compose.aws-postgres-test.yaml | 2 - docker-compose.azure-mysql-test.yaml | 2 - docker-compose.azure-postgres-test.yaml | 2 - docker-compose.basic-auth-test.yaml | 87 +++++++++++++++++++ docker-compose.gcp-mysql-test.yaml | 2 - docker-compose.gcp-postgres-test.yaml | 2 - test-containers/basic-auth/basic_auth.ini | 7 ++ tests/extended_docker_compose.py | 39 +++++++++ tests/test_basic_auth_postgres_integration.py | 74 ++++++++++++++++ 10 files changed, 207 insertions(+), 12 deletions(-) create mode 100644 docker-compose.basic-auth-test.yaml create mode 100644 test-containers/basic-auth/basic_auth.ini create mode 100644 tests/extended_docker_compose.py create mode 100644 tests/test_basic_auth_postgres_integration.py diff --git a/docker-compose.aws-mysql-test.yaml b/docker-compose.aws-mysql-test.yaml index 5bc610295..55a3a660f 100644 --- a/docker-compose.aws-mysql-test.yaml +++ b/docker-compose.aws-mysql-test.yaml @@ -1,5 +1,3 @@ -version: "3.8" - services: minio: container_name: minio-mysql-test diff --git a/docker-compose.aws-postgres-test.yaml b/docker-compose.aws-postgres-test.yaml index 8885d1fe9..5b634a531 100644 --- a/docker-compose.aws-postgres-test.yaml +++ b/docker-compose.aws-postgres-test.yaml @@ -1,5 +1,3 @@ -version: "3.8" - services: minio: container_name: minio-pg-test diff --git a/docker-compose.azure-mysql-test.yaml b/docker-compose.azure-mysql-test.yaml index fd9fb4305..da17240e9 100644 --- a/docker-compose.azure-mysql-test.yaml +++ b/docker-compose.azure-mysql-test.yaml @@ -1,5 +1,3 @@ -version: "3.8" - services: azurite: container_name: azurite-mysql-test diff --git a/docker-compose.azure-postgres-test.yaml b/docker-compose.azure-postgres-test.yaml index 6db5e617a..f57b25bda 100644 --- a/docker-compose.azure-postgres-test.yaml +++ b/docker-compose.azure-postgres-test.yaml @@ -1,5 +1,3 @@ -version: "3.8" - services: azurite: container_name: azurite-pg-test diff --git a/docker-compose.basic-auth-test.yaml b/docker-compose.basic-auth-test.yaml new file mode 100644 index 000000000..1383f7cda --- /dev/null +++ b/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 diff --git a/docker-compose.gcp-mysql-test.yaml b/docker-compose.gcp-mysql-test.yaml index 2ad3e4e0f..685e64f37 100644 --- a/docker-compose.gcp-mysql-test.yaml +++ b/docker-compose.gcp-mysql-test.yaml @@ -1,5 +1,3 @@ -version: "3.8" - services: gcs: container_name: gcs-mysql-test diff --git a/docker-compose.gcp-postgres-test.yaml b/docker-compose.gcp-postgres-test.yaml index a71e1894e..64cdc495a 100644 --- a/docker-compose.gcp-postgres-test.yaml +++ b/docker-compose.gcp-postgres-test.yaml @@ -1,5 +1,3 @@ -version: "3.8" - services: gcs: container_name: gcs-pg-test diff --git a/test-containers/basic-auth/basic_auth.ini b/test-containers/basic-auth/basic_auth.ini new file mode 100644 index 000000000..738768e04 --- /dev/null +++ b/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 diff --git a/tests/extended_docker_compose.py b/tests/extended_docker_compose.py new file mode 100644 index 000000000..7e0e23541 --- /dev/null +++ b/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." + ) diff --git a/tests/test_basic_auth_postgres_integration.py b/tests/test_basic_auth_postgres_integration.py new file mode 100644 index 000000000..5eff055b5 --- /dev/null +++ b/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"]