mirror of
https://github.com/datahub-project/datahub.git
synced 2025-11-03 12:16:10 +00:00
145 lines
4.2 KiB
Python
145 lines
4.2 KiB
Python
import json
|
|
|
|
import pytest
|
|
|
|
from datahub.ingestion.autogenerated.lineage_helper import (
|
|
_load_lineage_data,
|
|
clear_cache,
|
|
get_all_aspect_names,
|
|
get_lineage_data,
|
|
)
|
|
|
|
|
|
class TestLineageHelper:
|
|
@pytest.fixture
|
|
def mock_lineage_fields(self):
|
|
return [{"name": "dataset", "path": "upstreams.dataset", "isLineage": True}]
|
|
|
|
@pytest.fixture
|
|
def mock_lineage_data(self):
|
|
return {
|
|
"entities": {
|
|
"dataset": {
|
|
"upstreamLineage": {
|
|
"aspect": "upstreamLineage",
|
|
"fields": [
|
|
{
|
|
"name": "dataset",
|
|
"path": "upstreams.dataset",
|
|
"isLineage": True,
|
|
"relationship": {
|
|
"name": "DownstreamOf",
|
|
"entityTypes": ["dataset"],
|
|
"isLineage": True,
|
|
},
|
|
}
|
|
],
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
@pytest.fixture
|
|
def mock_file_data(self, mock_lineage_data):
|
|
return json.dumps(mock_lineage_data)
|
|
|
|
def setup_method(self):
|
|
clear_cache()
|
|
|
|
def teardown_method(self):
|
|
clear_cache()
|
|
|
|
def setup_mock_get_fields(self, monkeypatch, fields):
|
|
def mock_get_fields(*args, **kwargs):
|
|
return fields
|
|
|
|
monkeypatch.setattr(
|
|
"datahub.ingestion.autogenerated.lineage_helper.get_lineage_fields",
|
|
mock_get_fields,
|
|
)
|
|
|
|
def setup_mock_load_data(self, monkeypatch, data):
|
|
def mock_load_data():
|
|
return data
|
|
|
|
monkeypatch.setattr(
|
|
"datahub.ingestion.autogenerated.lineage_helper._load_lineage_data",
|
|
mock_load_data,
|
|
)
|
|
|
|
def setup_mock_file_operations(self, monkeypatch, file_data, exists=True):
|
|
def mock_open_file(*args, **kwargs):
|
|
class MockFile:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def read(self):
|
|
return file_data
|
|
|
|
return MockFile()
|
|
|
|
def mock_path_exists(*args, **kwargs):
|
|
return exists
|
|
|
|
monkeypatch.setattr("builtins.open", mock_open_file)
|
|
monkeypatch.setattr("pathlib.Path.exists", mock_path_exists)
|
|
|
|
def test_load_lineage_data_success(
|
|
self, monkeypatch, mock_file_data, mock_lineage_data
|
|
):
|
|
self.setup_mock_file_operations(monkeypatch, mock_file_data, exists=True)
|
|
|
|
result = _load_lineage_data()
|
|
|
|
assert result == mock_lineage_data
|
|
assert (
|
|
result["entities"]["dataset"]["upstreamLineage"]["fields"][0]["isLineage"]
|
|
is True
|
|
)
|
|
|
|
def test_load_lineage_data_file_not_found(self, monkeypatch):
|
|
self.setup_mock_file_operations(monkeypatch, "", exists=False)
|
|
|
|
# Should return empty dict instead of raising exception
|
|
result = _load_lineage_data()
|
|
assert result == {}
|
|
|
|
def test_load_lineage_data_invalid_json(self, monkeypatch):
|
|
self.setup_mock_file_operations(monkeypatch, "invalid json", exists=True)
|
|
|
|
# Should return empty dict instead of raising exception
|
|
result = _load_lineage_data()
|
|
assert result == {}
|
|
|
|
def test_get_all_aspect_names(self, monkeypatch, mock_lineage_data):
|
|
self.setup_mock_load_data(monkeypatch, mock_lineage_data)
|
|
|
|
clear_cache()
|
|
|
|
aspect_names = get_all_aspect_names()
|
|
|
|
expected_aspects = ["upstreamLineage"]
|
|
assert aspect_names == expected_aspects
|
|
|
|
def test_get_all_aspect_names_empty_entities(self, monkeypatch):
|
|
self.setup_mock_load_data(monkeypatch, {"entities": {}})
|
|
|
|
clear_cache()
|
|
|
|
aspect_names = get_all_aspect_names()
|
|
|
|
assert aspect_names == []
|
|
|
|
|
|
def test_get_all_lineage_aspect_names():
|
|
lineage_data = get_lineage_data()
|
|
entity_names = lineage_data.entities.keys()
|
|
assert "dataset" in entity_names
|
|
assert (
|
|
lineage_data.entities["dataset"].aspects["upstreamLineage"].fields[0].name
|
|
== "dataset"
|
|
)
|