File size: 4,961 Bytes
322f5cd 1bd23ce 322f5cd 0510038 322f5cd 1bd23ce 0510038 322f5cd 1bd23ce 0510038 322f5cd 1bd23ce 0510038 b304992 0510038 322f5cd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
import unittest
from app import (
render_mermaid,
extract_lineage_from_text,
extract_lineage_from_bigquery,
extract_lineage_from_url,
)
class TestLineageExtractors(unittest.TestCase):
def test_render_mermaid_wraps_and_inits(self):
viz = "graph TD\n A --> B"
html = render_mermaid(viz)
# Using mermaid.ink service for SVG rendering
self.assertIn('mermaid.ink/svg/', html)
self.assertIn('<img', html)
self.assertIn('mermaid.live/edit', html)
def test_extract_lineage_from_text_returns_html_and_summary(self):
# Test with valid JSON input
sample_json = '{"nodes": [{"id": "a", "name": "A"}], "edges": []}'
html, summary = extract_lineage_from_text(sample_json, "Custom JSON", "Mermaid")
self.assertIsInstance(html, str)
self.assertIsInstance(summary, str)
self.assertIn('mermaid.ink/svg/', html)
self.assertIn('Parsed', summary)
def test_extract_lineage_from_text_empty_input(self):
# Test with empty input
html, summary = extract_lineage_from_text("", "dbt Manifest", "Mermaid")
self.assertIsInstance(html, str)
self.assertIsInstance(summary, str)
self.assertIn('provide metadata', summary.lower())
def test_extract_lineage_from_bigquery_returns_html_and_summary(self):
html, summary = extract_lineage_from_bigquery("proj", "SELECT 1", "key", "Mermaid")
self.assertIn('mermaid.ink/svg/', html)
self.assertIn('BigQuery', summary)
def test_extract_lineage_from_url_returns_html_and_summary(self):
html, summary = extract_lineage_from_url("https://example.com", "Mermaid")
self.assertIn('mermaid.ink/svg/', html)
# Summary can be either 'Lineage' or 'Parsed' depending on response
self.assertTrue('Lineage' in summary or 'Parsed' in summary)
class TestExporters(unittest.TestCase):
def test_openlineage_export(self):
from exporters import LineageGraph, LineageNode, LineageEdge, OpenLineageExporter
graph = LineageGraph(name="test")
graph.add_node(LineageNode(id="a", name="Node A", type="table"))
graph.add_node(LineageNode(id="b", name="Node B", type="table"))
graph.add_edge(LineageEdge(source="a", target="b", type="transform"))
exporter = OpenLineageExporter(graph)
output = exporter.export()
self.assertIn("openlineage", output.lower())
self.assertIn("Node A", output)
def test_collibra_export(self):
from exporters import LineageGraph, LineageNode, LineageEdge, CollibraExporter
graph = LineageGraph(name="test")
graph.add_node(LineageNode(id="a", name="Node A", type="table"))
exporter = CollibraExporter(graph)
output = exporter.export()
self.assertIn("Collibra", output)
self.assertIn("Node A", output)
def test_purview_export(self):
from exporters import LineageGraph, LineageNode, LineageEdge, PurviewExporter
graph = LineageGraph(name="test")
graph.add_node(LineageNode(id="a", name="Node A", type="table"))
exporter = PurviewExporter(graph)
output = exporter.export()
self.assertIn("Purview", output)
self.assertIn("Node A", output)
def test_alation_export(self):
from exporters import LineageGraph, LineageNode, LineageEdge, AlationExporter
graph = LineageGraph(name="test")
graph.add_node(LineageNode(id="a", name="Node A", type="table"))
exporter = AlationExporter(graph)
output = exporter.export()
self.assertIn("Alation", output)
self.assertIn("Node A", output)
def test_atlas_export(self):
from exporters import LineageGraph, LineageNode, LineageEdge, AtlasExporter
graph = LineageGraph(name="test")
graph.add_node(LineageNode(id="a", name="Node A", type="table"))
graph.add_node(LineageNode(id="b", name="Node B", type="table"))
graph.add_edge(LineageEdge(source="a", target="b", type="transform"))
exporter = AtlasExporter(graph)
output = exporter.export()
self.assertIn("Atlas", output)
self.assertIn("Node A", output)
self.assertIn("entities", output)
class TestSampleDataLoading(unittest.TestCase):
def test_load_sample_simple(self):
from app import load_sample
content = load_sample("simple")
self.assertIn("nodes", content)
self.assertIn("edges", content)
def test_load_sample_dbt(self):
from app import load_sample
content = load_sample("dbt")
self.assertIn("metadata", content)
self.assertIn("nodes", content)
def test_load_sample_airflow(self):
from app import load_sample
content = load_sample("airflow")
self.assertIn("dag_id", content)
self.assertIn("tasks", content)
if __name__ == '__main__':
unittest.main()
|