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()