tests/test_app.py

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