mermaid-go/pkg/detector/detector_test.go

package detector

import (
	"testing"
)

func TestDetectDiagramType(t *testing.T) {
	registry := NewDetectorRegistry()

	tests := []struct {
		name     string
		input    string
		expected DiagramType
	}{
		{
			name:     "Sequence Diagram",
			input:    "sequenceDiagram\n    A -> B : hello",
			expected: DiagramTypeSequence,
		},
		{
			name:     "State Diagram V2",
			input:    "stateDiagram-v2\n    [*] --> State1",
			expected: DiagramTypeStateV2,
		},
		{
			name:     "State Diagram",
			input:    "stateDiagram\n    [*] --> State1",
			expected: DiagramTypeState,
		},
		{
			name:     "Flowchart",
			input:    "flowchart TD\n    A --> B",
			expected: DiagramTypeFlowchart,
		},
		{
			name:     "Graph (alias for flowchart)",
			input:    "graph TD\n    A --> B",
			expected: DiagramTypeFlowchart,
		},
		{
			name:     "Class Diagram",
			input:    "classDiagram\n    class Animal",
			expected: DiagramTypeClass,
		},
		{
			name:     "ER Diagram",
			input:    "erDiagram\n    CUSTOMER ||--o{ ORDER",
			expected: DiagramTypeER,
		},
		{
			name:     "Gantt Chart",
			input:    "gantt\n    title A Gantt Diagram",
			expected: DiagramTypeGantt,
		},
		{
			name:     "Pie Chart",
			input:    "pie\n    title Pets adopted by volunteers",
			expected: DiagramTypePie,
		},
		{
			name:     "Git Graph",
			input:    "gitGraph\n    commit",
			expected: DiagramTypeGit,
		},
		{
			name:     "User Journey",
			input:    "journey\n    title My working day",
			expected: DiagramTypeUserJourney,
		},
		{
			name:     "Timeline",
			input:    "timeline\n    title History of Social Media Platform",
			expected: DiagramTypeTimeline,
		},
		{
			name:     "Mindmap",
			input:    "mindmap\n  root((mindmap))",
			expected: DiagramTypeMindmap,
		},
		{
			name:     "Kanban",
			input:    "kanban\n    title Kanban Example",
			expected: DiagramTypeKanban,
		},
		{
			name:     "Sankey Diagram",
			input:    "sankey-beta\n    A --> B",
			expected: DiagramTypeSankey,
		},
		{
			name:     "XY Chart",
			input:    "xychart-beta\n    title Sales Revenue",
			expected: DiagramTypeXYChart,
		},
		{
			name:     "Quadrant Chart",
			input:    "quadrantChart\n    title Reach and engagement of campaigns",
			expected: DiagramTypeQuadrant,
		},
		{
			name:     "Requirement Diagram",
			input:    "requirement\n    requirement Test",
			expected: DiagramTypeRequirement,
		},
		{
			name:     "Block Diagram",
			input:    "block-beta\n    A",
			expected: DiagramTypeBlock,
		},
		{
			name:     "C4 Context",
			input:    "c4Context\n    title System Context diagram for Internet Banking System",
			expected: DiagramTypeC4,
		},
		{
			name:     "Architecture Diagram",
			input:    "architecture\n    title Architecture diagram",
			expected: DiagramTypeArchitecture,
		},
		{
			name:     "Radar Chart",
			input:    "radar\n    title Simple Radar Chart",
			expected: DiagramTypeRadar,
		},
		{
			name:     "Treemap",
			input:    "treemap\n    title Tree Map Example",
			expected: DiagramTypeTreemap,
		},
		{
			name:     "Packet Diagram",
			input:    "packet\n    title Packet Diagram",
			expected: DiagramTypePacket,
		},
		{
			name:     "Info Diagram",
			input:    "info\n    title Info Diagram",
			expected: DiagramTypeInfo,
		},
		{
			name:     "Unknown Diagram",
			input:    "unknownDiagram\n    some content",
			expected: DiagramTypeUnknown,
		},
		{
			name:     "Empty Input",
			input:    "",
			expected: DiagramTypeUnknown,
		},
		{
			name:     "With Comments",
			input:    "%% This is a comment\nsequenceDiagram\n    A -> B : hello",
			expected: DiagramTypeSequence,
		},
		{
			name:     "With Front Matter",
			input:    "---\ntitle: Test\n---\nsequenceDiagram\n    A -> B : hello",
			expected: DiagramTypeSequence,
		},
		{
			name:     "With Directives",
			input:    "%%{init: {'theme':'base'}}%%\nsequenceDiagram\n    A -> B : hello",
			expected: DiagramTypeSequence,
		},
		{
			name:     "Case Insensitive",
			input:    "SEQUENCEDIAGRAM\n    A -> B : hello",
			expected: DiagramTypeSequence,
		},
		{
			name:     "With Leading Whitespace",
			input:    "   sequenceDiagram\n    A -> B : hello",
			expected: DiagramTypeSequence,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := registry.DetectDiagramType(tt.input)
			if result != tt.expected {
				t.Errorf("Expected %v, got %v", tt.expected, result)
			}
		})
	}
}

func TestCleanText(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name:     "Remove Comments",
			input:    "%% This is a comment\nsequenceDiagram\n    A -> B : hello",
			expected: "sequenceDiagram\n    A -> B : hello",
		},
		{
			name:     "Remove Front Matter",
			input:    "---\ntitle: Test\n---\nsequenceDiagram\n    A -> B : hello",
			expected: "sequenceDiagram\n    A -> B : hello",
		},
		{
			name:     "Remove Directives",
			input:    "%%{init: {'theme':'base'}}%%\nsequenceDiagram\n    A -> B : hello",
			expected: "sequenceDiagram\n    A -> B : hello",
		},
		{
			name:     "Remove Multiple Comments",
			input:    "%% Comment 1\nsequenceDiagram\n    %% Comment 2\n    A -> B : hello",
			expected: "sequenceDiagram\n    A -> B : hello",
		},
		{
			name:     "Trim Whitespace",
			input:    "   \n  sequenceDiagram  \n  \n    A -> B : hello  \n  ",
			expected: "sequenceDiagram\n    A -> B : hello",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := cleanText(tt.input)
			if result != tt.expected {
				t.Errorf("Expected %q, got %q", tt.expected, result)
			}
		})
	}
}

func TestDetectorPriority(t *testing.T) {
	registry := NewDetectorRegistry()

	// Test that stateDiagram-v2 has higher priority than stateDiagram
	stateV2Input := "stateDiagram-v2\n    [*] --> State1"
	result := registry.DetectDiagramType(stateV2Input)
	if result != DiagramTypeStateV2 {
		t.Errorf("Expected %v, got %v", DiagramTypeStateV2, result)
	}

	// Test that sequenceDiagram has higher priority than generic patterns
	sequenceInput := "sequenceDiagram\n    A -> B : hello"
	result = registry.DetectDiagramType(sequenceInput)
	if result != DiagramTypeSequence {
		t.Errorf("Expected %v, got %v", DiagramTypeSequence, result)
	}
}