mermaid-go/pkg/parser/enhanced_test.go

package parser

import (
	"mermaid-go/pkg/renderer"
	"testing"
)

// TestEnhancedSequenceDiagramSupport tests the enhanced sequence diagram parsing and rendering
func TestEnhancedSequenceDiagramSupport(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name: "Basic sequence with different arrow types",
			input: `sequenceDiagram
    participant A
    participant B
    A -> B : solid arrow
    A --> B : dotted arrow
    A ->> B : solid arrow with arrowhead
    A -->> B : dotted arrow with arrowhead
    A -x B : destroy solid
    A --x B : destroy dotted
    A -) B : open solid
    A --) B : open dotted
    A <-> B : bidirectional`,
			expected: "sequenceDiagram",
		},
		{
			name: "Sequence with loops and alternatives",
			input: `sequenceDiagram
    participant U as User
    participant S as Server
    participant D as Database
    
    loop Authentication
        U -> S : login request
        S -> D : validate credentials
        D --> S : user data
    end
    
    alt Valid credentials
        S --> U : login success
    else Invalid credentials
        S --> U : login failed
    end
    
    opt Remember me
        S -> D : save session
    end`,
			expected: "sequenceDiagram",
		},
		{
			name: "Sequence with parallel sections",
			input: `sequenceDiagram
    participant A
    participant B
    participant C
    
    par Parallel execution
        A -> B : message 1
        A -> C : message 2
    and Another parallel section
        B -> C : message 3
    end`,
			expected: "sequenceDiagram",
		},
		{
			name: "Sequence with boxes and notes",
			input: `sequenceDiagram
    box "Frontend" lightblue
        participant U as User
        participant W as WebApp
    end
    
    box "Backend" lightgreen
        participant S as Server
        participant D as Database
    end
    
    U -> W : request
    W -> S : API call
    
    Note over W,S : Processing request
    
    S -> D : query
    D --> S : results
    S --> W : response
    W --> U : display`,
			expected: "sequenceDiagram",
		},
		{
			name: "Sequence with activations",
			input: `sequenceDiagram
    participant A
    participant B
    
    A -> B : start processing
    activate B
    B -> A : processing...
    deactivate B
    B --> A : done`,
			expected: "sequenceDiagram",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Parse the input
			parser := NewSequenceParser()
			diagram, err := parser.Parse(tt.input)
			if err != nil {
				t.Fatalf("Failed to parse: %v", err)
			}

			// Verify basic structure
			if diagram == nil {
				t.Fatal("Diagram is nil")
			}

			// Render back to mermaid syntax
			renderer := renderer.NewSequenceRenderer()
			output, err := renderer.Render(diagram)
			if err != nil {
				t.Fatalf("Failed to render: %v", err)
			}

			// Verify output contains expected elements
			if !containsString(output, tt.expected) {
				t.Errorf("Expected output to contain %q, got: %s", tt.expected, output)
			}

			t.Logf("Successfully parsed and rendered: %s", tt.name)
			t.Logf("Output: %s", output)
		})
	}
}

// TestEnhancedStateDiagramSupport tests the enhanced state diagram parsing and rendering
func TestEnhancedStateDiagramSupport(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name: "Basic state diagram with transitions",
			input: `stateDiagram-v2
    [*] --> State1
    State1 --> State2 : transition
    State2 --> [*]`,
			expected: "stateDiagram-v2",
		},
		{
			name: "State diagram with guards and actions",
			input: `stateDiagram-v2
    [*] --> Idle
    Idle --> Processing : start [hasData] / processData
    Processing --> Idle : complete [success] / cleanup
    Processing --> Error : fail [error] / logError`,
			expected: "stateDiagram-v2",
		},
		{
			name: "State diagram with composite states",
			input: `stateDiagram-v2
    [*] --> Active
    
    state Active {
        [*] --> Working
        Working --> Waiting
        Waiting --> Working
    }
    
    Active --> [*]`,
			expected: "stateDiagram-v2",
		},
		{
			name: "State diagram with notes",
			input: `stateDiagram-v2
    [*] --> State1
    State1 --> State2
    
    note right of State1 : This is State1
    note left of State2 : This is State2`,
			expected: "stateDiagram-v2",
		},
		{
			name: "State diagram with special states",
			input: `stateDiagram-v2
    [*] --> Start
    Start --> Choice : <<choice>>
    Choice --> Fork : <<fork>>
    Fork --> Join : <<join>>
    Join --> History : <<history>>
    History --> [*]`,
			expected: "stateDiagram-v2",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Parse the input
			parser := NewStateParser()
			diagram, err := parser.Parse(tt.input)
			if err != nil {
				t.Fatalf("Failed to parse: %v", err)
			}

			// Verify basic structure
			if diagram == nil {
				t.Fatal("Diagram is nil")
			}

			// Render back to mermaid syntax
			renderer := renderer.NewStateRenderer()
			output, err := renderer.Render(diagram)
			if err != nil {
				t.Fatalf("Failed to render: %v", err)
			}

			// Verify output contains expected elements
			if !containsString(output, tt.expected) {
				t.Errorf("Expected output to contain %q, got: %s", tt.expected, output)
			}

			t.Logf("Successfully parsed and rendered: %s", tt.name)
			t.Logf("Output: %s", output)
		})
	}
}

// TestRoundTripConversion tests that parsing and rendering preserves the original structure
func TestRoundTripConversion(t *testing.T) {
	// Test sequence diagram round trip
	seqInput := `sequenceDiagram
    participant A
    participant B
    A -> B : hello
    loop repeat
        B --> A : response
    end
    alt success
        A -> B : continue
    else error
        A -> B : retry
    end`

	seqParser := NewSequenceParser()
	seqDiagram, err := seqParser.Parse(seqInput)
	if err != nil {
		t.Fatalf("Failed to parse sequence diagram: %v", err)
	}

	seqRenderer := renderer.NewSequenceRenderer()
	seqOutput, err := seqRenderer.Render(seqDiagram)
	if err != nil {
		t.Fatalf("Failed to render sequence diagram: %v", err)
	}

	t.Logf("Sequence diagram round trip successful")
	t.Logf("Original: %s", seqInput)
	t.Logf("Rendered: %s", seqOutput)

	// Test state diagram round trip
	stateInput := `stateDiagram-v2
    [*] --> State1
    State1 --> State2 : transition [guard] / action
    State2 --> [*]`

	stateParser := NewStateParser()
	stateDiagram, err := stateParser.Parse(stateInput)
	if err != nil {
		t.Fatalf("Failed to parse state diagram: %v", err)
	}

	stateRenderer := renderer.NewStateRenderer()
	stateOutput, err := stateRenderer.Render(stateDiagram)
	if err != nil {
		t.Fatalf("Failed to render state diagram: %v", err)
	}

	t.Logf("State diagram round trip successful")
	t.Logf("Original: %s", stateInput)
	t.Logf("Rendered: %s", stateOutput)
}

// TestComplexSequenceDiagram tests a complex sequence diagram with all features
func TestComplexSequenceDiagram(t *testing.T) {
	complexInput := `sequenceDiagram
    participant U as User
    participant W as WebApp
    participant A as API
    participant D as Database
    participant C as Cache
    
    box "Frontend" lightblue
        U
        W
    end
    
    box "Backend" lightgreen
        A
        D
        C
    end
    
    U -> W : login request
    activate W
    
    W -> A : authenticate
    activate A
    
    alt Cache hit
        A -> C : get user data
        C --> A : cached data
    else Cache miss
        A -> D : query user
        D --> A : user data
        A -> C : store in cache
    end
    
    A --> W : auth result
    deactivate A
    
    opt Remember session
        W -> C : store session
    end
    
    W --> U : login response
    deactivate W
    
    Note over U,C : Authentication complete
    
    par User actions
        U -> W : browse content
        W -> A : fetch data
    and Background tasks
        A -> D : cleanup old sessions
    end`

	parser := NewSequenceParser()
	diagram, err := parser.Parse(complexInput)
	if err != nil {
		t.Fatalf("Failed to parse complex sequence diagram: %v", err)
	}

	// Verify structure
	if len(diagram.Participants) < 5 {
		t.Errorf("Expected at least 5 participants, got %d", len(diagram.Participants))
	}

	if len(diagram.Boxes) < 2 {
		t.Errorf("Expected at least 2 boxes, got %d", len(diagram.Boxes))
	}

	if len(diagram.Alts) < 1 {
		t.Errorf("Expected at least 1 alt block, got %d", len(diagram.Alts))
	}

	if len(diagram.Opts) < 1 {
		t.Errorf("Expected at least 1 opt block, got %d", len(diagram.Opts))
	}

	if len(diagram.Pars) < 1 {
		t.Errorf("Expected at least 1 par block, got %d", len(diagram.Pars))
	}

	if len(diagram.Notes) < 1 {
		t.Errorf("Expected at least 1 note, got %d", len(diagram.Notes))
	}

	// Render and verify
	renderer := renderer.NewSequenceRenderer()
	output, err := renderer.Render(diagram)
	if err != nil {
		t.Fatalf("Failed to render complex sequence diagram: %v", err)
	}

	t.Logf("Complex sequence diagram test passed")
	t.Logf("Output: %s", output)
}

// Helper function to check if a string contains a substring
func containsString(s, substr string) bool {
	return len(s) >= len(substr) &&
		(s == substr ||
			(len(s) > len(substr) &&
				(s[:len(substr)] == substr ||
					s[len(s)-len(substr):] == substr ||
					containsSubstring(s, substr))))
}

func containsSubstring(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}