open-metadata--openmetadata/ingestion/tests/integration/conftest.py

import logging
import time
from typing import List, Tuple, Type  # noqa: UP035

import pytest

from _openmetadata_testutils.ometa import int_admin_ometa
from metadata.generated.schema.entity.services.databaseService import DatabaseService
from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
    AutoClassificationConfigType,
)
from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
    DatabaseMetadataConfigType,
)
from metadata.generated.schema.metadataIngestion.workflow import LogLevels
from metadata.ingestion.api.common import Entity
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.workflow.ingestion import IngestionWorkflow


@pytest.fixture(scope="session", autouse=True)
def configure_logging():
    logging.getLogger("sqlfluff").setLevel(logging.CRITICAL)
    logging.getLogger("pytds").setLevel(logging.CRITICAL)


@pytest.fixture(scope="session")
def metadata():
    return int_admin_ometa()


def pytest_pycollect_makeitem(collector, name, obj):
    try:
        bases = [base.__name__ for base in obj.mro()]
        for cls in ("BaseModel", "Enum"):
            if cls in bases:
                return []
    except (AttributeError, TypeError):
        pass


@pytest.hookimpl(tryfirst=True)
def pytest_collection_modifyitems(config, items):
    """Auto-apply xdist_group markers based on source directory.

    This ensures all tests for a source (mysql, postgres, etc.) run in the
    SAME worker, allowing them to share module-scoped container fixtures.

    Examples:
        tests/integration/mysql/test_metadata.py      -> xdist_group("mysql_suite")
        tests/integration/postgres/test_metadata.py   -> xdist_group("postgres_suite")
        tests/integration/ometa/test_domains.py       -> xdist_group("ometa_suite")
    """
    for item in items:
        test_path = str(item.fspath)

        if "/integration/" in test_path:
            parts = test_path.split("/integration/")
            if len(parts) > 1:
                source_name = parts[1].split("/")[0]
                group_name = f"{source_name}_suite"
                item.add_marker(pytest.mark.xdist_group(group_name))


# TODO: Will be addressed when cleaning up integration tests.
#  Setting the max tries for testcontainers here has pitfalls,
#  the main one being that it cannot be changed through the recommended
#  way of using environment variables. The main problem is that
#  waiting_utils.py uses testcontainers_config.timeout as a default
#  value for the timeout. Therefore, if we want to effectively change
#  this value, we must do so before the module is imported,
#  which is a potential source of issues.


@pytest.fixture(scope="session", autouse=True)
def config_testcontatiners():
    try:
        from testcontainers.core.config import testcontainers_config
    except ModuleNotFoundError:

        class _TestContainerConfig:  # type: ignore[too-few-public-methods]
            max_tries = 10

        return _TestContainerConfig()

    testcontainers_config.max_tries = 10
    return testcontainers_config


@pytest.fixture(scope="session")
def sink_config(metadata):
    return {
        "type": "metadata-rest",
        "config": {},
    }


@pytest.fixture(scope="session")
def workflow_config(metadata):
    return {
        "loggerLevel": LogLevels.DEBUG.value,
        "openMetadataServerConfig": metadata.config.model_dump(),
    }


@pytest.fixture(scope="module")
def profiler_config(db_service, workflow_config, sink_config):
    return {
        "source": {
            "type": db_service.connection.config.type.value.lower(),
            "serviceName": db_service.fullyQualifiedName.root,
            "sourceConfig": {
                "config": {
                    "type": "Profiler",
                    "timeoutSeconds": 600,
                    "threadCount": 1,  # easier for debugging
                }
            },
        },
        "processor": {
            "type": "orm-profiler",
            "config": {},
        },
        "sink": sink_config,
        "workflowConfig": workflow_config,
    }


@pytest.fixture(scope="module")
def classifier_config(db_service, workflow_config, sink_config):
    return {
        "source": {
            "type": db_service.connection.config.type.value.lower(),
            "serviceName": db_service.fullyQualifiedName.root,
            "sourceConfig": {
                "config": {
                    "type": AutoClassificationConfigType.AutoClassification.value,
                    "storeSampleData": True,
                    "enableAutoClassification": True,
                }
            },
        },
        "processor": {
            "type": "orm-profiler",
            "config": {},
        },
        "sink": sink_config,
        "workflowConfig": workflow_config,
    }


@pytest.fixture(scope="module")
def run_workflow():
    def _run(workflow_type: Type[IngestionWorkflow], config, raise_from_status=True):  # noqa: UP006
        workflow: IngestionWorkflow = workflow_type.create(config)
        workflow.execute()
        if raise_from_status:
            workflow.print_status()
            workflow.raise_from_status()
        return workflow

    return _run


logger = logging.getLogger(__name__)


def _safe_delete(metadata, entity, entity_id, retries=3, **kwargs):
    """Delete with retry logic to handle transient server errors during parallel teardown.

    A 404 here means the entity is already gone (e.g., wiped as part of an earlier
    cascade or another worker's teardown); treat it as success rather than retrying.
    """
    for attempt in range(retries):
        try:
            metadata.delete(entity=entity, entity_id=entity_id, **kwargs)
            return  # noqa: TRY300
        except Exception as exc:
            if _is_not_found(exc):
                logger.debug("Skipping %s %s delete — already gone", entity.__name__, entity_id)
                return
            if attempt < retries - 1:
                logger.warning(
                    "Retry %d/%d: delete %s %s",
                    attempt + 1,
                    retries,
                    entity.__name__,
                    entity_id,
                )
                time.sleep(0.5 * (attempt + 1))
            else:
                raise


def _is_not_found(exc: BaseException) -> bool:
    status = getattr(getattr(exc, "response", None), "status_code", None)
    if status == 404:
        return True
    return "404" in str(exc)


@pytest.fixture(scope="module")
def db_service(metadata, create_service_request, unmask_password):
    service_entity = metadata.create_or_update(data=create_service_request)
    fqn = service_entity.fullyQualifiedName.root
    yield unmask_password(service_entity)
    service_entity = metadata.get_by_name(DatabaseService, fqn)
    if service_entity:
        _safe_delete(
            metadata,
            entity=DatabaseService,
            entity_id=service_entity.id,
            recursive=True,
            hard_delete=True,
        )


@pytest.fixture(scope="module")
def unmask_password(create_service_request):
    """Unmask the db passwrod returned by the metadata service.
    You can override this at the test_module level to implement custom password handling.

    Example:
    @pytest.fixture(scope="module")
    def unmask_password(my_container1, my_container2):
        def patch_password(service: DatabaseService):
            if service.connection.config.authType.password == "my_password":
              ... # do something else
            return service
        return patch_password
    """

    def patch_password(service: DatabaseService):
        if hasattr(service.connection.config, "authType"):
            service.connection.config.authType.password = create_service_request.connection.config.authType.password
            return service
        service.connection.config.password = create_service_request.connection.config.password
        return service

    return patch_password


@pytest.fixture(scope="module")
def create_service_request():
    """
    Implement in the test module to create a service request
    Example:
    def create_service_request(scope="module"):
        return CreateDatabaseServiceRequest(
            name="my_service",
            serviceType=DatabaseServiceType.MyService,
            connection=DatabaseConnection(
                config=MyServiceConnection(
                    username="my_user",
                    password="my_password",
                    host="localhost",
                    port="5432",
                )
            ),
        )
    """
    raise NotImplementedError("Implement in the test module")


@pytest.fixture(scope="module")
def monkeymodule():
    with pytest.MonkeyPatch.context() as mp:
        yield mp


@pytest.fixture(scope="module")
def patch_passwords_for_db_services(db_service, unmask_password, monkeymodule):
    """Patch the password for all db services returned by the metadata service.

    Usage:

    def test_my_test(db_service, patch_passwords_for_db_services):
        ...

    OR

    @pytest.usefixtures("patch_passwords_for_db_services")
    def test_my_test(db_service):
        ...
    """

    def override_password(getter):
        def inner(*args, **kwargs):
            result = getter(*args, **kwargs)
            if isinstance(result, DatabaseService):  # noqa: SIM102
                if result.fullyQualifiedName.root == db_service.fullyQualifiedName.root:
                    return unmask_password(result)
            return result

        return inner

    monkeymodule.setattr(
        "metadata.ingestion.ometa.ometa_api.OpenMetadata.get_by_name",
        override_password(OpenMetadata.get_by_name),
    )

    monkeymodule.setattr(
        "metadata.ingestion.ometa.ometa_api.OpenMetadata.get_by_id",
        override_password(OpenMetadata.get_by_id),
    )


@pytest.fixture
def cleanup_fqns(metadata):
    fqns: List[Tuple[Type[Entity], str]] = []  # noqa: UP006

    def inner(entity_type: Type[Entity], fqn: str):  # noqa: UP006
        fqns.append((entity_type, fqn))

    yield inner
    for etype, fqn in fqns:
        entity = metadata.get_by_name(etype, fqn, fields=["*"])
        if entity:
            _safe_delete(
                metadata,
                entity=etype,
                entity_id=entity.id,
                recursive=True,
                hard_delete=True,
            )


@pytest.fixture(scope="module")
def ingestion_config(db_service, metadata, workflow_config, sink_config):
    return {
        "source": {
            "type": db_service.connection.config.type.value.lower(),
            "serviceName": db_service.fullyQualifiedName.root,
            "sourceConfig": {"config": {"type": DatabaseMetadataConfigType.DatabaseMetadata.value}},
            "serviceConnection": db_service.connection.model_dump(),
        },
        "sink": sink_config,
        "workflowConfig": workflow_config,
    }