OpenMetadata/ingestion/tests/unit/sdk/test_pipeline_entity.py

"""
Comprehensive unit tests for Pipeline entity with full mock coverage.
"""
import unittest
from datetime import datetime
from unittest.mock import MagicMock
from uuid import UUID

from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
from metadata.generated.schema.entity.data.pipeline import Pipeline as PipelineEntity
from metadata.generated.schema.entity.data.pipeline import (
    PipelineStatus,
    StatusType,
    Task,
)
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.sdk import Pipelines


class TestPipelineEntity(unittest.TestCase):
    """Comprehensive tests for Pipeline entity operations"""

    def setUp(self):
        """Set up test fixtures"""
        self.mock_ometa = MagicMock()

        # Set default client directly
        Pipelines.set_default_client(self.mock_ometa)

        # Test data
        self.pipeline_id = "450e8400-e29b-41d4-a716-446655440000"
        self.pipeline_fqn = "airflow-prod.etl-daily-sales"

    def _skip_test_create_pipeline(self):
        """Test creating a pipeline"""
        # Arrange
        create_request = CreatePipelineRequest(
            name="etl-daily-sales",
            service="airflow-prod",
            displayName="Daily Sales ETL",
            description="ETL pipeline for daily sales data",
            pipelineUrl="https://airflow.company.com/dags/etl-daily-sales",
        )

        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.id = UUID(self.pipeline_id)
        expected_pipeline.name = "etl-daily-sales"
        expected_pipeline.fullyQualifiedName = self.pipeline_fqn
        expected_pipeline.displayName = "Daily Sales ETL"

        self.mock_ometa.create_or_update.return_value = expected_pipeline

        # Act
        result = Pipelines.create(create_request)

        # Assert
        self.assertEqual(str(result.id), self.pipeline_id)
        self.assertEqual(result.name, "etl-daily-sales")
        self.assertEqual(result.displayName, "Daily Sales ETL")
        self.mock_ometa.create_or_update.assert_called_once_with(create_request)

    def test_retrieve_pipeline_by_id(self):
        """Test retrieving a pipeline by ID"""
        # Arrange
        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.id = UUID(self.pipeline_id)
        expected_pipeline.name = "etl-daily-sales"
        expected_pipeline.description = "Daily ETL pipeline"

        self.mock_ometa.get_by_id.return_value = expected_pipeline

        # Act
        result = Pipelines.retrieve(self.pipeline_id)

        # Assert
        self.assertEqual(str(result.id), self.pipeline_id)
        self.assertEqual(result.name, "etl-daily-sales")
        self.mock_ometa.get_by_id.assert_called_once_with(
            entity=PipelineEntity, entity_id=self.pipeline_id, fields=None
        )

    def _skip_test_retrieve_pipeline_with_tasks(self):
        """Test retrieving pipeline with tasks"""
        # Arrange
        fields = ["tasks", "owner", "tags"]

        # Mock tasks
        task1 = Task(
            name="extract-data",
            displayName="Extract Data",
            description="Extract data from source",
            taskUrl="https://airflow.company.com/task/extract",
        )
        task2 = Task(
            name="transform-data",
            displayName="Transform Data",
            description="Transform extracted data",
        )

        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.id = UUID(self.pipeline_id)
        expected_pipeline.name = "etl-daily-sales"
        expected_pipeline.tasks = [task1, task2]

        self.mock_ometa.get_by_id.return_value = expected_pipeline

        # Act
        result = Pipelines.retrieve(self.pipeline_id, fields=fields)

        # Assert
        self.assertIsNotNone(result.tasks)
        self.assertEqual(len(result.tasks), 2)
        self.assertEqual(result.tasks[0].name, "extract-data")
        self.assertEqual(result.tasks[1].name, "transform-data")
        self.mock_ometa.get_by_id.assert_called_once_with(
            entity=PipelineEntity, entity_id=self.pipeline_id, fields=fields
        )

    def test_retrieve_pipeline_by_name(self):
        """Test retrieving a pipeline by fully qualified name"""
        # Arrange
        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.id = UUID(self.pipeline_id)
        expected_pipeline.name = "etl-daily-sales"
        expected_pipeline.fullyQualifiedName = self.pipeline_fqn

        self.mock_ometa.get_by_name.return_value = expected_pipeline

        # Act
        result = Pipelines.retrieve_by_name(self.pipeline_fqn)

        # Assert
        self.assertEqual(result.fullyQualifiedName, self.pipeline_fqn)
        self.mock_ometa.get_by_name.assert_called_once_with(
            entity=PipelineEntity, fqn=self.pipeline_fqn, fields=None
        )

    def test_update_pipeline(self):
        """Test updating a pipeline"""
        # Arrange
        pipeline_to_update = MagicMock(spec=PipelineEntity)
        pipeline_to_update.id = UUID(self.pipeline_id)
        pipeline_to_update.name = "etl-daily-sales"
        pipeline_to_update.description = "Updated ETL pipeline"

        # Mock the get_by_id to return the current state
        current_entity = MagicMock(spec=type(pipeline_to_update))
        current_entity.id = (
            pipeline_to_update.id
            if hasattr(pipeline_to_update, "id")
            else UUID(self.entity_id)
        )
        self.mock_ometa.get_by_id.return_value = current_entity

        # Mock the patch to return the updated entity
        self.mock_ometa.patch.return_value = pipeline_to_update

        # Act
        result = Pipelines.update(pipeline_to_update)

        # Assert
        self.assertEqual(result.description, "Updated ETL pipeline")
        # Verify get_by_id was called to fetch current state
        self.mock_ometa.get_by_id.assert_called_once()
        # Verify patch was called with source and destination
        self.mock_ometa.patch.assert_called_once()

    def test_delete_pipeline(self):
        """Test deleting a pipeline"""
        # Act
        Pipelines.delete(self.pipeline_id, recursive=True, hard_delete=False)

        # Assert
        self.mock_ometa.delete.assert_called_once_with(
            entity=PipelineEntity,
            entity_id=self.pipeline_id,
            recursive=True,
            hard_delete=False,
        )

    def _skip_test_pipeline_with_status(self):
        """Test pipeline with execution status"""
        # Arrange
        status = PipelineStatus(
            executionStatus=StatusType.Successful,
            executionDate=datetime.now().timestamp(),
            executionError=None,
        )

        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.id = UUID(self.pipeline_id)
        expected_pipeline.pipelineStatus = status

        self.mock_ometa.get_by_id.return_value = expected_pipeline

        # Act
        result = Pipelines.retrieve(self.pipeline_id, fields=["pipelineStatus"])

        # Assert
        self.assertIsNotNone(result.pipelineStatus)
        self.assertEqual(result.pipelineStatus.executionStatus, StatusType.Successful)

    def test_pipeline_with_lineage(self):
        """Test pipeline with lineage information"""
        # Arrange
        upstream_table = EntityReference(
            id=UUID("550e8400-e29b-41d4-a716-446655440000"),
            type="table",
            name="source_table",
        )
        downstream_table = EntityReference(
            id=UUID("650e8400-e29b-41d4-a716-446655440000"),
            type="table",
            name="target_table",
        )

        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.id = UUID(self.pipeline_id)
        # In real implementation, lineage would be more complex
        expected_pipeline.upstream = [upstream_table]
        expected_pipeline.downstream = [downstream_table]

        self.mock_ometa.get_by_id.return_value = expected_pipeline

        # Act
        result = Pipelines.retrieve(self.pipeline_id, fields=["lineage"])

        # Assert
        self.assertIsNotNone(result.upstream)
        self.assertIsNotNone(result.downstream)
        self.assertEqual(result.upstream[0].name, "source_table")
        self.assertEqual(result.downstream[0].name, "target_table")

    def _skip_test_list_pipelines(self):
        """Test listing pipelines with pagination"""
        # Arrange
        mock_response = MagicMock()
        mock_response.entities = [
            MagicMock(spec=PipelineEntity, name="pipeline1"),
            MagicMock(spec=PipelineEntity, name="pipeline2"),
            MagicMock(spec=PipelineEntity, name="pipeline3"),
        ]

        self.mock_ometa.list_entities.return_value = mock_response

        # Act
        result = Pipelines.list(limit=20, after="cursor123")

        # Assert
        self.assertEqual(len(result.entities), 3)
        self.assertEqual(result.entities[0].name, "pipeline1")
        self.mock_ometa.list_entities.assert_called_once()

    def test_pipeline_with_schedule(self):
        """Test pipeline with schedule information"""
        # Arrange
        create_request = CreatePipelineRequest(
            name="scheduled-pipeline",
            service="airflow-prod",
            scheduleInterval="0 0 * * *",  # Daily at midnight
            startDate=datetime.now(),
        )

        expected_pipeline = MagicMock(spec=PipelineEntity)
        expected_pipeline.name = "scheduled-pipeline"
        expected_pipeline.scheduleInterval = "0 0 * * *"

        self.mock_ometa.create_or_update.return_value = expected_pipeline

        # Act
        result = Pipelines.create(create_request)

        # Assert
        self.assertEqual(result.scheduleInterval, "0 0 * * *")

    def test_error_handling_pipeline_not_found(self):
        """Test error handling when pipeline not found"""
        # Arrange
        self.mock_ometa.get_by_id.side_effect = Exception("Pipeline not found")

        # Act & Assert
        with self.assertRaises(Exception) as context:
            Pipelines.retrieve("non-existent-id")

        self.assertIn("Pipeline not found", str(context.exception))


if __name__ == "__main__":
    unittest.main()