diff --git a/docker-compose.basic-auth-test.yaml b/docker-compose.basic-auth-postgres-test.yaml similarity index 96% rename from docker-compose.basic-auth-test.yaml rename to docker-compose.basic-auth-postgres-test.yaml index 1383f7cda..bb6472400 100644 --- a/docker-compose.basic-auth-test.yaml +++ b/docker-compose.basic-auth-postgres-test.yaml @@ -66,7 +66,7 @@ services: - basic_auth_test_nw volumes: - mlflow-basic-auth-storage:/mlflow - - ./test-containers/basic-auth/basic_auth.ini:/mlflow/basic_auth.ini + - ./test-containers/basic-auth/postgres/basic_auth.ini:/mlflow/basic_auth.ini depends_on: - minio - postgres diff --git a/test-containers/basic-auth/basic_auth.ini b/test-containers/basic-auth/postgres/basic_auth.ini similarity index 100% rename from test-containers/basic-auth/basic_auth.ini rename to test-containers/basic-auth/postgres/basic_auth.ini diff --git a/tests/extended_docker_compose.py b/tests/extended_docker_compose.py index 7e0e23541..7de9bfed1 100644 --- a/tests/extended_docker_compose.py +++ b/tests/extended_docker_compose.py @@ -1,39 +1,54 @@ import re -import subprocess import time +import docker 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 + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.docker_client = docker.from_env() - 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] + def wait_for_logs(self, service_name, expected_log, timeout=120, interval=5): + """ + Wait for a specific log entry in the Docker service's logs. + :param service_name: The name of the service in the Docker Compose file. + :param expected_log: A string or regex pattern to match in the service logs. + :param timeout: Maximum time to wait for the log in seconds. + :param interval: Time interval between log checks in seconds. + """ 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}") + logs = self.get_service_logs(service_name) + + # Use regex to match the expected log pattern + if re.search(expected_log, logs): + print(f"Found expected log matching: {expected_log}") + return + time.sleep(interval) + raise TimeoutError( f"Log message matching '{expected_log!r}' not found within {timeout} seconds." ) + + def get_service_logs(self, service_name): + """ + Fetch logs for a specific service using the Docker SDK. + + :param service_name: The name of the service in the Docker Compose file. + :return: The logs as a decoded string. + """ + try: + containers = self.docker_client.containers.list(filters={"name": service_name}) + if not containers: + raise ValueError(f"No container found for service '{service_name!r}'") + + logs = containers[0].logs().decode("utf-8") + return logs + except Exception as e: + # Explicitly re-raise the exception with additional context + raise RuntimeError(f"Error fetching logs for service '{service_name!r}': {e}") from e diff --git a/tests/test_basic_auth_postgres_integration.py b/tests/test_basic_auth_postgres_integration.py index 3d4af357b..3e7cf3cac 100644 --- a/tests/test_basic_auth_postgres_integration.py +++ b/tests/test_basic_auth_postgres_integration.py @@ -15,7 +15,7 @@ def test_postgres_backended_model_upload_and_access_with_basic_auth( ): with ExtendedDockerCompose( context=".", - compose_file_name=["docker-compose.basic-auth-test.yaml"], + compose_file_name=["docker-compose.basic-auth-postgres-test.yaml"], # pull=True, build=True, ) as compose: