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

52 lines
1.2 KiB
Python

from dataclasses import dataclass
import pytest
from datahub.ingestion.source.datahub.datahub_database_reader import (
VersionOrderable,
VersionOrderer,
)
@dataclass
class MockRow(VersionOrderable):
createdon: int
version: int
urn: str
@pytest.fixture
def rows():
return [
MockRow(0, 0, "one"),
MockRow(0, 1, "one"),
MockRow(0, 0, "two"),
MockRow(0, 0, "three"),
MockRow(0, 1, "three"),
MockRow(0, 2, "three"),
MockRow(0, 1, "two"),
MockRow(0, 4, "three"),
MockRow(0, 5, "three"),
MockRow(1, 6, "three"),
MockRow(1, 0, "four"),
MockRow(2, 0, "five"),
MockRow(2, 1, "six"),
MockRow(2, 0, "six"),
MockRow(3, 0, "seven"),
MockRow(3, 0, "eight"),
]
def test_version_orderer(rows):
orderer = VersionOrderer[MockRow](enabled=True)
ordered_rows = list(orderer(rows))
assert ordered_rows == sorted(
ordered_rows, key=lambda x: (x.createdon, x.version == 0)
)
def test_version_orderer_disabled(rows):
orderer = VersionOrderer[MockRow](enabled=False)
ordered_rows = list(orderer(rows))
assert ordered_rows == rows