fix(ingest/ci): fix docker compose integration test failures (#15424)

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

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

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

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)