mirror of
https://github.com/open-metadata/OpenMetadata.git
synced 2025-07-08 01:29:37 +00:00
84 lines
2.2 KiB
Python
84 lines
2.2 KiB
Python
|
# 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.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,
|
||
|
|
context="numbers",
|
||
|
|
processor="yield_numbers",
|
||
|
|
ack_sink=False,
|
||
|
|
)
|
||
|
|
],
|
||
|
|
children=["strings"],
|
||
|
|
)
|
||
|
|
strings = TopologyNode(
|
||
|
|
producer="get_strings",
|
||
|
|
stages=[
|
||
|
|
NodeStage(
|
||
|
|
type_=str,
|
||
|
|
context="strings",
|
||
|
|
processor="yield_strings",
|
||
|
|
ack_sink=False,
|
||
|
|
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 number + 1
|
||
|
|
|
||
|
|
def yield_strings(self, my_str: str):
|
||
|
|
yield my_str + str(self.context.numbers)
|
||
|
|
|
||
|
|
|
||
|
|
class TopologyRunnerTest(TestCase):
|
||
|
|
"""
|
||
|
|
Validate filter patterns
|
||
|
|
"""
|
||
|
|
|
||
|
|
def test_node_and_stage(self):
|
||
|
|
source = MockSource()
|
||
|
|
processed = list(source.next_record())
|
||
|
|
assert processed == [2, "abc2", "def2", 3, "abc3", "def3"]
|