OpenMetadata/ingestion/tests/unit/topology/test_runner.py

#  Copyright 2021 Collate
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  http://www.apache.org/licenses/LICENSE-2.0
#  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.

"""
Check that we are properly running nodes and stages
"""
from unittest import TestCase

from metadata.ingestion.api.models import Either
from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
from metadata.ingestion.models.topology import (
    NodeStage,
    ServiceTopology,
    TopologyNode,
    create_source_context,
)


class MockTopology(ServiceTopology):
    root = TopologyNode(
        producer="get_numbers",
        stages=[
            NodeStage(
                type_=int,
                processor="yield_numbers",
            )
        ],
        children=["strings"],
    )
    strings = TopologyNode(
        producer="get_strings",
        stages=[
            NodeStage(
                type_=str,
                processor="yield_strings",
                consumer=["numbers"],
            )
        ],
    )


class MockSource(TopologyRunnerMixin):
    topology = MockTopology()
    context = create_source_context(topology)

    @staticmethod
    def get_numbers():
        yield 1
        yield 2

    @staticmethod
    def get_strings():
        yield "abc"
        yield "def"

    @staticmethod
    def yield_numbers(number: int):
        yield Either(right=number + 1)

    @staticmethod
    def yield_strings(my_str: str):
        yield Either(right=my_str)


class TopologyRunnerTest(TestCase):
    """Validate filter patterns"""

    @staticmethod
    def test_node_and_stage():
        source = MockSource()
        processed = list(source._iter())
        assert [either.right for either in processed] == [
            2,
            "abc",
            "def",
            3,
            "abc",
            "def",
        ]