datahub/metadata-ingestion/tests/unit/cli/docker/test_docker_check.py

import pathlib
from unittest.mock import MagicMock

import pytest

from datahub.cli.docker_check import (
    MIN_DISK_SPACE_NEEDED,
    MIN_MEMORY_NEEDED,
    ContainerStatus,
    DockerContainerStatus,
    DockerLowDiskSpaceError,
    DockerLowMemoryError,
    QuickstartStatus,
    check_upgrade_supported,
    run_quickstart_preflight_checks,
)


@pytest.fixture
def mock_docker_client():
    client = MagicMock()
    return client


def test_run_quickstart_preflight_checks_success(mock_docker_client):
    # Mock Docker info to return sufficient memory
    mock_docker_client.info.return_value = {
        "MemTotal": int(MIN_MEMORY_NEEDED * 1024 * 1024 * 1000)  # Convert GB to bytes
    }

    # Mock container run to return sufficient disk space
    mock_container = MagicMock()
    mock_container.decode.return_value = f"{MIN_DISK_SPACE_NEEDED * 1024**3} {MIN_DISK_SPACE_NEEDED * 1024**3}"  # total, available in bytes
    mock_docker_client.containers.run.return_value = mock_container

    # Test should pass without raising any exceptions
    run_quickstart_preflight_checks(mock_docker_client)


def test_run_quickstart_preflight_checks_insufficient_memory(mock_docker_client):
    # Mock Docker info to return insufficient memory
    mock_docker_client.info.return_value = {
        "MemTotal": int(
            (MIN_MEMORY_NEEDED - 1) * 1024 * 1024 * 1000
        )  # 1GB less than required
    }

    # Test should raise DockerLowMemoryError
    with pytest.raises(DockerLowMemoryError) as exc_info:
        run_quickstart_preflight_checks(mock_docker_client)

    assert "Total Docker memory configured" in str(exc_info.value)


def test_run_quickstart_preflight_checks_insufficient_disk_space(mock_docker_client):
    # Mock Docker info to return sufficient memory
    mock_docker_client.info.return_value = {
        "MemTotal": int(MIN_MEMORY_NEEDED * 1024 * 1024 * 1000)
    }

    # Mock container run to return insufficient disk space
    mock_container = MagicMock()
    mock_container.decode.return_value = f"{MIN_DISK_SPACE_NEEDED * 1024**3} {(MIN_DISK_SPACE_NEEDED - 1) * 1024**3}"  # total, available in bytes
    mock_docker_client.containers.run.return_value = mock_container

    # Test should raise DockerLowDiskSpaceError
    with pytest.raises(DockerLowDiskSpaceError) as exc_info:
        run_quickstart_preflight_checks(mock_docker_client)

    assert "Total Docker disk space available" in str(exc_info.value)


@pytest.fixture
def mock_compose_file(tmp_path):
    compose_file = tmp_path / "docker-compose.yml"
    compose_file.write_text("""
services:
  datahub-gms:
    image: datahub-gms
  datahub-frontend:
    image: datahub-frontend
volumes:
  datahub_volume1:
  datahub_volume2:
    """)
    return compose_file


def test_check_upgrade_supported_no_containers(mock_compose_file):
    # Test when no containers are running
    quickstart_status = QuickstartStatus([], [], running_unsupported_version=False)
    assert check_upgrade_supported([mock_compose_file], quickstart_status) is True


def test_check_upgrade_supported_legacy_version():
    # Test when running legacy version
    quickstart_status = QuickstartStatus(
        [DockerContainerStatus("zookeeper", ContainerStatus.OK)],
        [],
        running_unsupported_version=True,
    )
    assert (
        check_upgrade_supported([pathlib.Path("dummy.yml")], quickstart_status) is False
    )


def test_check_upgrade_supported_matching_services(mock_compose_file):
    # Test when services and volumes match
    quickstart_status = QuickstartStatus(
        [
            DockerContainerStatus("datahub-gms", ContainerStatus.OK),
            DockerContainerStatus("datahub-frontend", ContainerStatus.OK),
        ],
        ["datahub_volume1", "datahub_volume2"],
        running_unsupported_version=False,
    )
    assert check_upgrade_supported([mock_compose_file], quickstart_status) is True


def test_check_upgrade_supported_mismatched_services(mock_compose_file):
    # Test when services don't match
    quickstart_status = QuickstartStatus(
        [
            DockerContainerStatus("datahub-gms", ContainerStatus.OK),
            DockerContainerStatus("datahub-frontend", ContainerStatus.OK),
            DockerContainerStatus(
                "zookeeper", ContainerStatus.OK
            ),  # Extra service not in compose
        ],
        ["datahub_volume1", "datahub_volume2"],
        running_unsupported_version=False,
    )
    assert check_upgrade_supported([mock_compose_file], quickstart_status) is False


def test_check_upgrade_supported_mismatched_volumes(mock_compose_file):
    # Test when volumes don't match
    quickstart_status = QuickstartStatus(
        [
            DockerContainerStatus("datahub-gms", ContainerStatus.OK),
            DockerContainerStatus("datahub-frontend", ContainerStatus.OK),
        ],
        [
            "datahub_volume1",
            "datahub_volume2",
            "zookeeper_volume",
        ],  # Extra volume not in compose
        running_unsupported_version=False,
    )
    assert check_upgrade_supported([mock_compose_file], quickstart_status) is False