From 6b8b6aaa08521bedeae4044c161b1736dce9571d Mon Sep 17 00:00:00 2001
From: kyungsoo-datahub <kyungsoo.lee@datahub.com>
Date: Wed, 26 Nov 2025 11:39:18 -0800
Subject: [PATCH] fix(ingest/ci): fix docker compose integration test failures
 (#15424)

---
 .../src/datahub/testing/docker_utils.py       |  3 ++-
 .../integration/hex/docker/docker-compose.yml | 24 ++-----------------
 .../kafka-connect/test_kafka_connect.py       |  4 +++-
 .../tests/integration/trino/test_trino.py     |  2 +-
 4 files changed, 8 insertions(+), 25 deletions(-)

diff --git a/metadata-ingestion/src/datahub/testing/docker_utils.py b/metadata-ingestion/src/datahub/testing/docker_utils.py
index 3501c3910c..1a7b457eaa 100644
--- a/metadata-ingestion/src/datahub/testing/docker_utils.py
+++ b/metadata-ingestion/src/datahub/testing/docker_utils.py
@@ -61,6 +61,7 @@ def docker_compose_runner(
         key: str,
         cleanup: bool = True,
         parallel: int = DOCKER_DEFAULT_UNLIMITED_PARALLELISM,
+        setup_command: Optional[Union[List[str], str]] = None,
     ) -> Iterator[pytest_docker.plugin.Services]:
         with pytest_docker.plugin.get_docker_services(
             docker_compose_command=f"{docker_compose_command} --parallel {parallel}",
@@ -68,7 +69,7 @@ def docker_compose_runner(
             # https://github.com/avast/pytest-docker/pull/108
             docker_compose_file=compose_file_path,  # type: ignore
             docker_compose_project_name=f"{docker_compose_project_name}-{key}",
-            docker_setup=docker_setup,
+            docker_setup=setup_command if setup_command is not None else docker_setup,
             docker_cleanup=docker_cleanup if cleanup else [],
         ) as docker_services:
             yield docker_services
diff --git a/metadata-ingestion/tests/integration/hex/docker/docker-compose.yml b/metadata-ingestion/tests/integration/hex/docker/docker-compose.yml
index 78eb0ae494..b63ed9b7f2 100644
--- a/metadata-ingestion/tests/integration/hex/docker/docker-compose.yml
+++ b/metadata-ingestion/tests/integration/hex/docker/docker-compose.yml
@@ -12,17 +12,7 @@ services:
       - -c
       - |
         apk add --no-cache wget
-        python /app/mock_hex_server.py &
-        SERVER_PID=$$!
-        for i in $$(seq 1 30); do
-          if wget --no-verbose --tries=1 --spider http://localhost:8000/health 2>/dev/null; then
-            wait $$SERVER_PID
-            exit 0
-          fi
-          sleep 1
-        done
-        kill $$SERVER_PID 2>/dev/null || true
-        exit 1
+        python /app/mock_hex_server.py
     healthcheck:
       test: ["CMD-SHELL", "wget --no-verbose --tries=1 --spider http://localhost:8000/health || exit 1"]
       interval: 5s
@@ -45,17 +35,7 @@ services:
       - -c
       - |
         apk add --no-cache wget
-        python /app/mock_datahub_server.py &
-        SERVER_PID=$$!
-        for i in $$(seq 1 30); do
-          if wget --no-verbose --tries=1 --spider http://localhost:8010/health 2>/dev/null; then
-            wait $$SERVER_PID
-            exit 0
-          fi
-          sleep 1
-        done
-        kill $$SERVER_PID 2>/dev/null || true
-        exit 1
+        python /app/mock_datahub_server.py
     healthcheck:
       test: ["CMD-SHELL", "wget --no-verbose --tries=1 --spider http://localhost:8010/health || exit 1"]
       interval: 5s
diff --git a/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py b/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py
index 4a7470d0f9..748a5451c2 100644
--- a/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py
+++ b/metadata-ingestion/tests/integration/kafka-connect/test_kafka_connect.py
@@ -128,7 +128,9 @@ def kafka_connect_runner(docker_compose_runner, pytestconfig, test_resources_dir
         str(test_resources_dir / "docker-compose.override.yml"),
     ]
 
-    with docker_compose_runner(docker_compose_file, "kafka-connect") as docker_services:
+    with docker_compose_runner(
+        docker_compose_file, "kafka-connect", setup_command=["up --wait"]
+    ) as docker_services:
         # We rely on Docker health checks to confirm all services are up & healthy
 
         # However healthcheck for test_connect service is not very trustable, so
diff --git a/metadata-ingestion/tests/integration/trino/test_trino.py b/metadata-ingestion/tests/integration/trino/test_trino.py
index e70da281d5..2cabdda228 100644
--- a/metadata-ingestion/tests/integration/trino/test_trino.py
+++ b/metadata-ingestion/tests/integration/trino/test_trino.py
@@ -29,7 +29,7 @@ data_platform = "trino"
 def trino_runner(docker_compose_runner, pytestconfig):
     test_resources_dir = pytestconfig.rootpath / "tests/integration/trino"
     with docker_compose_runner(
-        test_resources_dir / "docker-compose.yml", "trino"
+        test_resources_dir / "docker-compose.yml", "trino", setup_command=["up --wait"]
     ) as docker_services:
         wait_for_port(docker_services, "testtrino", 8080)
         wait_for_port(docker_services, "testhiveserver2", 10000, timeout=120)