mermaid-go/pkg/parser/enhanced_state_test.go

package parser

import (
	"strings"
	"testing"

	"mermaid-go/pkg/ast"
	"mermaid-go/pkg/renderer"
)

func TestStateParser_EnhancedFeatures(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			name: "Basic State with Actions",
			input: `stateDiagram-v2
    [*] --> State1
    State1 : entry / initialize
    State1 : exit / cleanup
    State1 : do / process
    State1 --> State2 : transition`,
			expected: `stateDiagram-v2
    [*] --> State1
    State1 : entry / initialize
    State1 : exit / cleanup
    State1 : do / process
    State1 --> State2 : transition
`,
		},
		{
			name: "Composite State",
			input: `stateDiagram-v2
    [*] --> CompositeState
    state CompositeState {
        [*] --> SubState1
        SubState1 --> SubState2
        SubState2 --> [*]
    }
    CompositeState --> [*]`,
			expected: `stateDiagram-v2
    [*] --> CompositeState
    state CompositeState {
        [*] --> SubState1
        SubState1 --> SubState2
        SubState2 --> [*]
    }
    CompositeState --> [*]
`,
		},
		{
			name: "Parallel States",
			input: `stateDiagram-v2
    [*] --> ParallelState
    state ParallelState {
        [*] --> Region1
        [*] --> Region2
        state Region1 {
            [*] --> StateA
            StateA --> StateB
        }
        state Region2 {
            [*] --> StateC
            StateC --> StateD
        }
    }`,
			expected: `stateDiagram-v2
    [*] --> ParallelState
    state ParallelState {
        [*] --> Region1
        [*] --> Region2
        state Region1 {
            [*] --> StateA
            StateA --> StateB
        }
        state Region2 {
            [*] --> StateC
            StateC --> StateD
        }
    }
`,
		},
		{
			name: "Special States",
			input: `stateDiagram-v2
    [*] --> ForkState
    ForkState : <<fork>>
    ForkState --> ChoiceState
    ChoiceState : <<choice>>
    ChoiceState --> JoinState
    JoinState : <<join>>
    JoinState --> [*]`,
			expected: `stateDiagram-v2
    state JoinState : <<join>>
    state ForkState : <<fork>>
    state ChoiceState : <<choice>>
    [*] --> ForkState
    ForkState --> ChoiceState
    ChoiceState --> JoinState
    JoinState --> [*]
`,
		},
		{
			name: "Complex Transitions with Guards and Actions",
			input: `stateDiagram-v2
    [*] --> State1
    State1 --> State2 : normal [condition] / action
    State2 --> State3 : urgent [priority > 5] / handle
    State3 --> [*]`,
			expected: `stateDiagram-v2
    [*] --> State1
    State1 --> State2 : normal [condition] / action
    State2 --> State3 : urgent [priority > 5] / handle
    State3 --> [*]
`,
		},
	}

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

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

			// Normalize whitespace for comparison
			expected := strings.TrimSpace(tt.expected)
			actual := strings.TrimSpace(output)

			if expected != actual {
				t.Errorf("Expected:\n%s\n\nGot:\n%s", expected, actual)
			}
		})
	}
}

func TestStateParser_StateActions(t *testing.T) {
	input := `stateDiagram-v2
    [*] --> State1
    State1 : entry / initialize system
    State1 : exit / cleanup resources
    State1 : do / process data
    State1 --> State2`

	parser := NewStateParser()
	diagram, err := parser.Parse(input)
	if err != nil {
		t.Fatalf("Failed to parse: %v", err)
	}

	// Check state was parsed with actions
	if len(diagram.States) == 0 {
		t.Fatal("No states found")
	}

	var state1 *ast.StateNode
	for _, state := range diagram.States {
		if state.ID == "State1" {
			state1 = state
			break
		}
	}
	if state1 == nil {
		t.Fatal("State1 not found")
	}

	if state1.EntryAction == nil || *state1.EntryAction != "/ initialize system" {
		t.Errorf("Expected entry action '/ initialize system', got %v", state1.EntryAction)
	}

	if state1.ExitAction == nil || *state1.ExitAction != "/ cleanup resources" {
		t.Errorf("Expected exit action '/ cleanup resources', got %v", state1.ExitAction)
	}

	if state1.DoAction == nil || *state1.DoAction != "/ process data" {
		t.Errorf("Expected do action '/ process data', got %v", state1.DoAction)
	}
}

func TestStateParser_CompositeStates(t *testing.T) {
	input := `stateDiagram-v2
    [*] --> CompositeState
    state CompositeState {
        [*] --> SubState1
        SubState1 --> SubState2
    }
    CompositeState --> [*]`

	parser := NewStateParser()
	diagram, err := parser.Parse(input)
	if err != nil {
		t.Fatalf("Failed to parse: %v", err)
	}

	// Check composite state was parsed
	var compositeState *ast.StateNode
	for _, state := range diagram.States {
		if state.ID == "CompositeState" {
			compositeState = state
			break
		}
	}
	if compositeState == nil {
		t.Fatal("CompositeState not found")
	}

	if compositeState.SubStates == nil || len(compositeState.SubStates) == 0 {
		t.Fatal("No substates found in CompositeState")
	}

	// Check substates
	if _, exists := compositeState.SubStates["SubState1"]; !exists {
		t.Error("SubState1 not found in CompositeState")
	}

	if _, exists := compositeState.SubStates["SubState2"]; !exists {
		t.Error("SubState2 not found in CompositeState")
	}
}