open-metadata--openmetadata/ingestion/tests/unit/sampler/test_container_sampler_processor.py

#  Copyright 2025 Collate
#  Licensed under the Collate Community License, Version 1.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
Test Container sampler processor functionality
"""

import uuid
from unittest.mock import MagicMock, Mock, patch

import pytest

from metadata.generated.schema.entity.data.container import (
    Container,
    ContainerDataModel,
)
from metadata.generated.schema.entity.data.table import (
    Column,
    ColumnName,
    DataType,
    Table,
    TableData,
)
from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
    OpenMetadataConnection,
)
from metadata.generated.schema.metadataIngestion.storageServiceAutoClassificationPipeline import (
    StorageServiceAutoClassificationPipeline,
)
from metadata.generated.schema.metadataIngestion.workflow import (
    OpenMetadataWorkflowConfig,
    Processor,
    Sink,
    Source,
    SourceConfig,
    WorkflowConfig,
)
from metadata.generated.schema.type.basic import FullyQualifiedEntityName, Uuid
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.profiler.source.model import ProfilerSourceAndEntity
from metadata.sampler.processor import SamplerProcessor


@pytest.fixture
def container_entity():
    """Create a test Container entity"""
    return Container(
        id=uuid.uuid4(),
        name="test_container",
        fullyQualifiedName=FullyQualifiedEntityName(root="s3_service.test_container"),
        service=EntityReference(
            id=Uuid(root=uuid.uuid4()),
            type="storageService",
            name="s3_service",
            fullyQualifiedName="s3_service",
        ),
        dataModel=ContainerDataModel(
            columns=[
                Column(name="id", dataType=DataType.INT),
                Column(name="name", dataType=DataType.STRING),
                Column(name="email", dataType=DataType.STRING),
            ]
        ),
    )


@pytest.fixture
def table_entity():
    """Create a test Table entity for comparison"""
    return Table(
        id=uuid.uuid4(),
        name="test_table",
        fullyQualifiedName=FullyQualifiedEntityName("mysql.db.test_table"),
        columns=[
            Column(name="id", dataType=DataType.INT),
            Column(name="name", dataType=DataType.STRING),
        ],
    )


@pytest.fixture
def workflow_config():
    """Create test workflow configuration"""
    config = OpenMetadataWorkflowConfig(
        source=Source(
            type="s3",
            serviceName="s3_service",
            sourceConfig=SourceConfig(
                config=StorageServiceAutoClassificationPipeline(storeSampleData=True, sampleDataCount=50),
            ),
        ),
        processor=Processor(type="orm-profiler", config={}),
        sink=Sink(type="metadata-rest", config={}),
        workflowConfig=WorkflowConfig(
            openMetadataServerConfig=OpenMetadataConnection(
                hostPort="localhost:8585/api",
            )
        ),
    )
    # Mock the serviceConnection structure
    config.source.serviceConnection = Mock()
    config.source.serviceConnection.root = Mock()
    config.source.serviceConnection.root.config = {}
    return config


@patch("metadata.sampler.processor.import_sampler_class")
def test_sampler_processor_handles_container(mock_import_sampler, container_entity, workflow_config):
    """Test that SamplerProcessor can handle Container entities"""

    mock_sampler_class = MagicMock()
    mock_sampler_instance = MagicMock()
    mock_sampler_instance.generate_sample_data.return_value = TableData(
        columns=[
            ColumnName(root="id"),
            ColumnName(root="name"),
            ColumnName(root="email"),
        ],
        rows=[
            ["1", "Alice", "alice@example.com"],
            ["2", "Bob", "bob@example.com"],
        ],
    )
    mock_sampler_class.create.return_value = mock_sampler_instance
    mock_import_sampler.return_value = mock_sampler_class

    metadata_mock = MagicMock()
    metadata_mock.get_profiler_config_settings.return_value = None

    processor = SamplerProcessor(
        config=workflow_config,
        metadata=metadata_mock,
    )

    profiler_source = MagicMock()
    record = ProfilerSourceAndEntity.model_construct(profiler_source=profiler_source, entity=container_entity)

    result = processor._run(record)

    assert result.right is not None
    assert result.left is None
    assert result.right.entity == container_entity
    assert result.right.sample_data is not None
    assert result.right.sample_data.store is True


@patch("metadata.sampler.processor.import_sampler_class")
def test_sampler_processor_handles_table(mock_import_sampler, table_entity, workflow_config):
    """Test that SamplerProcessor still handles Table entities correctly"""

    mock_sampler_class = MagicMock()
    mock_sampler_instance = MagicMock()
    mock_sampler_instance.generate_sample_data.return_value = TableData(
        columns=[
            ColumnName(root="id"),
            ColumnName(root="name"),
        ],
        rows=[
            ["1", "Alice"],
            ["2", "Bob"],
        ],
    )
    mock_sampler_class.create.return_value = mock_sampler_instance
    mock_import_sampler.return_value = mock_sampler_class

    metadata_mock = MagicMock()
    metadata_mock.get_profiler_config_settings.return_value = None

    with patch("metadata.utils.profiler_utils.get_context_entities") as mock_get_context:
        mock_database_entity = MagicMock()
        mock_get_context.return_value = (None, mock_database_entity, None)

        with patch("metadata.sampler.entity_adapters.build_database_service_conn_config") as mock_build_conn:
            mock_build_conn.return_value = {}

            with patch("metadata.sampler.entity_adapters.get_profile_sample_config") as mock_sample_cfg:
                from metadata.sampler.models import SampleConfig

                mock_sample_cfg.return_value = SampleConfig()

                with patch("metadata.sampler.entity_adapters.get_sample_data_count_config") as mock_count:
                    mock_count.return_value = 50

                    processor = SamplerProcessor(
                        config=workflow_config,
                        metadata=metadata_mock,
                    )

                    profiler_source = MagicMock()
                    record = ProfilerSourceAndEntity.model_construct(
                        profiler_source=profiler_source, entity=table_entity
                    )

                    result = processor._run(record)

                    assert result.right is not None
                    assert result.left is None
                    assert result.right.entity == table_entity


def test_sampler_processor_container_no_context_entities_needed(container_entity, workflow_config):
    """Test that container sampling doesn't require database/schema context"""

    with patch("metadata.sampler.processor.import_sampler_class") as mock_import:
        mock_sampler_class = MagicMock()
        mock_sampler_instance = MagicMock()
        mock_sampler_instance.generate_sample_data.return_value = TableData(columns=[], rows=[])
        mock_sampler_class.create.return_value = mock_sampler_instance
        mock_import.return_value = mock_sampler_class

        metadata_mock = MagicMock()
        metadata_mock.get_profiler_config_settings.return_value = None

        processor = SamplerProcessor(
            config=workflow_config,
            metadata=metadata_mock,
        )

        profiler_source = MagicMock()
        record = ProfilerSourceAndEntity.model_construct(profiler_source=profiler_source, entity=container_entity)

        processor._run(record)

        call_args = mock_sampler_class.create.call_args
        assert "schema_entity" not in call_args.kwargs
        assert "database_entity" not in call_args.kwargs
        assert call_args.kwargs["entity"] == container_entity


def test_sampler_processor_unsupported_entity_type(workflow_config):
    """Test that processor rejects unsupported entity types"""

    unsupported_entity = MagicMock()
    unsupported_entity.fullyQualifiedName.root = "unsupported.entity"

    with patch("metadata.sampler.processor.import_sampler_class"):
        metadata_mock = MagicMock()
        metadata_mock.get_profiler_config_settings.return_value = None

        processor = SamplerProcessor(
            config=workflow_config,
            metadata=metadata_mock,
        )

        profiler_source = MagicMock()
        record = ProfilerSourceAndEntity.model_construct(profiler_source=profiler_source, entity=unsupported_entity)

        result = processor._run(record)

        assert result.left is not None
        assert result.right is None
        assert "Unsupported entity type" in result.left.error


def test_sample_data_store_flag_respected(container_entity, workflow_config):
    """Test that storeSampleData flag is properly passed to SampleData"""

    workflow_config.source.sourceConfig.config.storeSampleData = False

    with patch("metadata.sampler.processor.import_sampler_class") as mock_import:
        mock_sampler_class = MagicMock()
        mock_sampler_instance = MagicMock()
        mock_sampler_instance.generate_sample_data.return_value = TableData(columns=[], rows=[])
        mock_sampler_class.create.return_value = mock_sampler_instance
        mock_import.return_value = mock_sampler_class

        metadata_mock = MagicMock()
        metadata_mock.get_profiler_config_settings.return_value = None

        processor = SamplerProcessor(
            config=workflow_config,
            metadata=metadata_mock,
        )

        profiler_source = MagicMock()
        record = ProfilerSourceAndEntity.model_construct(profiler_source=profiler_source, entity=container_entity)

        result = processor._run(record)

        assert result.right.sample_data.store is False