added tests, removed mermaid dependency - lineage extractor

app.py

@@ -8,6 +8,27 @@ import json
 import os
 from typing import Optional, Tuple
 
+
+def render_mermaid(viz_code: str) -> str:
+    """Wrap mermaid source in HTML and initialize mermaid when the HTML is inserted.
+
+    This function also escapes angle brackets in the source to avoid HTML injection
+    while keeping the mermaid syntax intact.
+    """
+    # Escape HTML-sensitive characters
+    safe_viz = viz_code.replace("<", "&lt;").replace(">", "&gt;")
+    # Script will wait for mermaid to be available then initialize diagrams.
+    init_script = (
+        "<script>"
+        "(function(){"
+        "function run(){"
+        "  if(window.mermaid){ mermaid.init(undefined, document.querySelectorAll('.mermaid')); }"
+        "  else { setTimeout(run,50); }"
+        " } run();})();"
+        "</script>"
+    )
+    return f"<div class=\"mermaid\">{safe_viz}</div>{init_script}"
+
 # Note: This is a template. You'll need to integrate with your actual agent backend.
 # This could be through an API, Claude SDK, or other agent framework.
 
@@ -30,8 +51,9 @@ def extract_lineage_from_text(
     # TODO: Integrate with your agent backend
     # This is where you'd call your agent with the metadata_parser and graph_visualizer workers
     
+    viz = "graph TD\n    A[Sample Node] --> B[Output Node]"
     return (
-        "graph TD\n    A[Sample Node] --> B[Output Node]",
+        render_mermaid(viz),
         f"Processed {source_type} metadata. Found X nodes and Y relationships."
     )
 
@@ -55,8 +77,9 @@ def extract_lineage_from_bigquery(
     """
     # TODO: Integrate with BigQuery and your agent backend
     
+    viz = "graph TD\n    A[BigQuery Table] --> B[Destination Table]"
     return (
-        "graph TD\n    A[BigQuery Table] --> B[Destination Table]",
+        render_mermaid(viz),
         f"Extracted lineage from BigQuery project: {project_id}"
     )
 
@@ -76,8 +99,9 @@ def extract_lineage_from_url(
     """
     # TODO: Integrate with URL fetching and your agent backend
     
+    viz = "graph TD\n    A[API Source] --> B[Data Pipeline]"
     return (
-        "graph TD\n    A[API Source] --> B[Data Pipeline]",
+        render_mermaid(viz),
         f"Extracted lineage from URL: {url}"
     )
 
@@ -95,6 +119,12 @@ with gr.Blocks(title="Lineage Graph Extractor", theme=gr.themes.Soft()) as demo:
     - **URLs/APIs**: Fetch metadata from web endpoints
     - **dbt, Airflow, Snowflake**: Through MCP integration (when configured)
     """)
+    # Load Mermaid.js once (hidden). We set startOnLoad=false and will initialize
+    # individual diagrams after inserting them into the DOM.
+    mermaid_loader = gr.HTML(
+        value='''<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>\n<script>mermaid.initialize({startOnLoad:false});</script>''',
+        visible=False
+    )
     
     with gr.Tabs():
         # Tab 1: Text/File Input
@@ -119,9 +149,9 @@ with gr.Blocks(title="Lineage Graph Extractor", theme=gr.themes.Soft()) as demo:
                     extract_btn_text = gr.Button("Extract Lineage", variant="primary")
                 
                 with gr.Column():
-                    output_viz_text = gr.Code(
-                        label="Lineage Visualization",
-                        language="mermaid"
+                    output_viz_text = gr.HTML(
+                        value="",
+                        label="Lineage Visualization"
                     )
                     output_summary_text = gr.Textbox(
                         label="Summary",
@@ -160,9 +190,9 @@ with gr.Blocks(title="Lineage Graph Extractor", theme=gr.themes.Soft()) as demo:
                     extract_btn_bq = gr.Button("Extract Lineage", variant="primary")
                 
                 with gr.Column():
-                    output_viz_bq = gr.Code(
-                        label="Lineage Visualization",
-                        language="mermaid"
+                    output_viz_bq = gr.HTML(
+                        value="",
+                        label="Lineage Visualization"
                     )
                     output_summary_bq = gr.Textbox(
                         label="Summary",
@@ -191,9 +221,9 @@ with gr.Blocks(title="Lineage Graph Extractor", theme=gr.themes.Soft()) as demo:
                     extract_btn_url = gr.Button("Extract Lineage", variant="primary")
                 
                 with gr.Column():
-                    output_viz_url = gr.Code(
-                        label="Lineage Visualization",
-                        language="mermaid"
+                    output_viz_url = gr.HTML(
+                        value="",
+                        label="Lineage Visualization"
                     )
                     output_summary_url = gr.Textbox(
                         label="Summary",

tests/test_app.py

@@ -0,0 +1,38 @@
+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)
+        self.assertIn('<div class="mermaid">', html)
+        self.assertIn('graph TD', html)
+        self.assertIn('mermaid.init', html)
+
+    def test_extract_lineage_from_text_returns_html_and_summary(self):
+        html, summary = extract_lineage_from_text("", "dbt Manifest", "Mermaid")
+        self.assertIsInstance(html, str)
+        self.assertIsInstance(summary, str)
+        self.assertIn('<div class="mermaid">', html)
+        self.assertIn('Processed', summary)
+
+    def test_extract_lineage_from_bigquery_returns_html_and_summary(self):
+        html, summary = extract_lineage_from_bigquery("proj", "SELECT 1", "key", "Mermaid")
+        self.assertIn('<div class="mermaid">', html)
+        self.assertIn('Extracted lineage', summary)
+
+    def test_extract_lineage_from_url_returns_html_and_summary(self):
+        html, summary = extract_lineage_from_url("https://example.com", "Mermaid")
+        self.assertIn('<div class="mermaid">', html)
+        self.assertIn('Extracted lineage from URL', summary)
+
+
+if __name__ == '__main__':
+    unittest.main()