chore: import upstream snapshot with attribution

ingestion/tests/unit/sampler/test_container_sampler_processor.py

@@ -0,0 +1,288 @@
+#  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