datahub/metadata-ingestion/tests/unit/test_pipeline.py

import unittest
from unittest.mock import patch

from gometa.ingestion.run.pipeline import Pipeline


class PipelineTest(unittest.TestCase):
    @patch("gometa.ingestion.source.kafka.KafkaSource.get_workunits")
    @patch("gometa.ingestion.sink.console.ConsoleSink.close")
    def test_configure(self, mock_sink, mock_source):
        pipeline = Pipeline.create(
            {
                "source": {"type": "kafka", "config": {"bootstrap": "localhost:9092"}},
                "sink": {"type": "console"},
            }
        )
        pipeline.run()
        mock_source.assert_called_once()
        mock_sink.assert_called_once()