From 4d2edc6ef4fc8022240b32c4ea9312235cf9a939 Mon Sep 17 00:00:00 2001
From: Burak Ince <burak.ince@linux.org.tr>
Date: Sun, 15 Dec 2024 13:29:17 +0100
Subject: [PATCH] Add second wait for auth migrations

---
 docker-compose.basic-auth-postgres-test.yaml  | 26 +++++++++----------
 tests/test_basic_auth_postgres_integration.py |  7 ++---
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/docker-compose.basic-auth-postgres-test.yaml b/docker-compose.basic-auth-postgres-test.yaml
index bb6472400..98ad6899a 100644
--- a/docker-compose.basic-auth-postgres-test.yaml
+++ b/docker-compose.basic-auth-postgres-test.yaml
@@ -1,6 +1,6 @@
 services:
   minio:
-    container_name: minio-basic-auth-test
+    container_name: minio-basic-auth-pg-test
     hostname: minio
     image: minio/minio:${MINIO_VERSION}
     entrypoint: sh
@@ -12,13 +12,13 @@ services:
       - "9000:9000"
       - "9001:9001"
     volumes:
-      - minio-basic-auth-storage:/data
+      - minio-basic-auth-pg-storage:/data
     networks:
-      - basic_auth_test_nw
+      - basic_auth_pg_test_nw
 
   postgres:
     image: "postgres:${POSTGRES_VERSION}"
-    container_name: mlflow-basic-auth-db
+    container_name: mlflow-basic-auth-pg-db
     environment:
       POSTGRES_PASSWORD: postgres
       POSTGRES_USER: postgres
@@ -28,11 +28,11 @@ services:
     volumes:
       - postgres-basic-auth-storage:/var/lib/postgresql/data
     networks:
-      - basic_auth_test_nw
+      - basic_auth_pg_test_nw
 
   usersdb:
     image: "postgres:${POSTGRES_VERSION}"
-    container_name: mlflow-basic-auth-users-db
+    container_name: mlflow-basic-auth-users-pg-db
     environment:
       POSTGRES_PASSWORD: postgres
       POSTGRES_USER: postgres
@@ -42,10 +42,10 @@ services:
     volumes:
       - postgres-basic-auth-users-storage:/var/lib/postgresql/data
     networks:
-      - basic_auth_test_nw
+      - basic_auth_pg_test_nw
 
   mlflow:
-    container_name: mlflow-basic-auth-test
+    container_name: mlflow-basic-auth-pg-test
     build:
       context: .
       dockerfile: Dockerfile
@@ -63,9 +63,9 @@ services:
     ports:
       - "8080:8080"
     networks:
-      - basic_auth_test_nw
+      - basic_auth_pg_test_nw
     volumes:
-      - mlflow-basic-auth-storage:/mlflow
+      - mlflow-basic-auth-pg-storage:/mlflow
       - ./test-containers/basic-auth/postgres/basic_auth.ini:/mlflow/basic_auth.ini
     depends_on:
       - minio
@@ -77,11 +77,11 @@ volumes:
     driver: local
   postgres-basic-auth-users-storage:
     driver: local
-  mlflow-basic-auth-storage:
+  mlflow-basic-auth-pg-storage:
     driver: local
-  minio-basic-auth-storage:
+  minio-basic-auth-pg-storage:
     driver: local
 
 networks:
-  basic_auth_test_nw:
+  basic_auth_pg_test_nw:
     driver: bridge
diff --git a/tests/test_basic_auth_postgres_integration.py b/tests/test_basic_auth_postgres_integration.py
index e20cc12d4..39be0b7f8 100644
--- a/tests/test_basic_auth_postgres_integration.py
+++ b/tests/test_basic_auth_postgres_integration.py
@@ -35,10 +35,11 @@ def test_postgres_backended_model_upload_and_access_with_basic_auth(
 
         log_message = f".*Listening at: {re.escape(base_url)}.*"
         compose.wait_for_logs("mlflow", log_message)
-        time.sleep(3)  # Wait 3 seconds more the get flask ready
+        compose.wait_for_logs("mlflow", ".*8606fa83a998, initial_migration")
+        time.sleep(5)  # Wait 5 seconds more the get flask ready
 
-        experiment_name = "aws-cloud-postgres-experiment"
-        model_name = "test-aws-pg-model"
+        experiment_name = "basic-auth-postgres-experiment"
+        model_name = "test-basic-auth-pg-model"
         stage_name = "Staging"
         os.environ["MLFLOW_TRACKING_USERNAME"] = mlflow_admin_username
         os.environ["MLFLOW_TRACKING_PASSWORD"] = mlflow_admin_password