From d7610dda8fac9277413915a20849c59a969f9900 Mon Sep 17 00:00:00 2001
From: Sidharth Vinod <sidharthv96@gmail.com>
Date: Mon, 12 Dec 2022 23:51:46 +0530
Subject: [PATCH] fix: Throw correct errors when parsing diagrams with errors

---
 cSpell.json                          |  4 ++
 packages/mermaid/src/Diagram.ts      |  2 +-
 packages/mermaid/src/diagram.spec.ts | 67 ++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 packages/mermaid/src/diagram.spec.ts

diff --git a/cSpell.json b/cSpell.json
index 45e7fe99a..5711d29dc 100644
--- a/cSpell.json
+++ b/cSpell.json
@@ -14,6 +14,7 @@
     "bilkent",
     "bisheng",
     "braintree",
+    "brkt",
     "brolin",
     "brotli",
     "classdef",
@@ -60,12 +61,14 @@
     "mindmaps",
     "mitigations",
     "mkdocs",
+    "mult",
     "orlandoni",
     "phpbb",
     "plantuml",
     "playfair",
     "pnpm",
     "podlite",
+    "quence",
     "ranksep",
     "rect",
     "rects",
@@ -81,6 +84,7 @@
     "substate",
     "sveidqvist",
     "techn",
+    "teststr",
     "treemap",
     "ts-nocheck",
     "tuleap",
diff --git a/packages/mermaid/src/Diagram.ts b/packages/mermaid/src/Diagram.ts
index a2349c255..83412e4aa 100644
--- a/packages/mermaid/src/Diagram.ts
+++ b/packages/mermaid/src/Diagram.ts
@@ -102,7 +102,6 @@ export const getDiagramFromText = (
   try {
     // Trying to find the diagram
     getDiagram(type);
-    return new Diagram(txt, parseError);
   } catch (error) {
     const loader = getDiagramLoader(type);
     if (!loader) {
@@ -118,6 +117,7 @@ export const getDiagramFromText = (
       return new Diagram(txt, parseError);
     });
   }
+  return new Diagram(txt, parseError);
 };
 
 export default Diagram;
diff --git a/packages/mermaid/src/diagram.spec.ts b/packages/mermaid/src/diagram.spec.ts
new file mode 100644
index 000000000..ebe088a86
--- /dev/null
+++ b/packages/mermaid/src/diagram.spec.ts
@@ -0,0 +1,67 @@
+import { describe, test, expect } from 'vitest';
+import Diagram, { getDiagramFromText } from './Diagram';
+import { addDetector } from './diagram-api/detectType';
+import { addDiagrams } from './diagram-api/diagram-orchestration';
+
+addDiagrams();
+
+describe('diagram detection', () => {
+  test('should detect inbuilt diagrams', () => {
+    const graph = getDiagramFromText('graph TD; A-->B') as Diagram;
+    expect(graph).toBeInstanceOf(Diagram);
+    expect(graph.type).toBe('flowchart-v2');
+    const sequence = getDiagramFromText(
+      'sequenceDiagram; Alice->>+John: Hello John, how are you?'
+    ) as Diagram;
+    expect(sequence).toBeInstanceOf(Diagram);
+    expect(sequence.type).toBe('sequence');
+  });
+
+  test('should detect external diagrams', async () => {
+    addDetector(
+      'loki',
+      (str) => str.startsWith('loki'),
+      () =>
+        Promise.resolve({
+          id: 'loki',
+          diagram: {
+            db: {},
+            parser: {
+              parse: () => {
+                // no-op
+              },
+              parser: {
+                yy: {},
+              },
+            },
+            renderer: {},
+            styles: {},
+          },
+        })
+    );
+    const diagram = (await getDiagramFromText('loki TD; A-->B')) as Diagram;
+    expect(diagram).toBeInstanceOf(Diagram);
+    expect(diagram.type).toBe('loki');
+  });
+
+  test('should throw the right error for incorrect diagram', () => {
+    expect(() => getDiagramFromText('graph TD; A-->')).toThrowErrorMatchingInlineSnapshot(`
+"Parse error on line 3:
+graph TD; A-->
+--------------^
+Expecting 'AMP', 'ALPHA', 'COLON', 'PIPE', 'TESTSTR', 'DOWN', 'DEFAULT', 'NUM', 'COMMA', 'MINUS', 'BRKT', 'DOT', 'PUNCTUATION', 'UNICODE_TEXT', 'PLUS', 'EQUALS', 'MULT', 'UNDERSCORE', got 'EOF'"
+		`);
+    expect(() => getDiagramFromText('sequenceDiagram; A-->B')).toThrowErrorMatchingInlineSnapshot(`
+"Parse error on line 1:
+...quenceDiagram; A-->B
+-----------------------^
+Expecting 'TXT', got 'NEWLINE'"
+		`);
+  });
+
+  test('should throw the right error for unregistered diagrams', () => {
+    expect(() => getDiagramFromText('thor TD; A-->B')).toThrowError(
+      'No diagram type detected for text: thor TD; A-->B'
+    );
+  });
+});