ai tests + bug fix on last with 2args + strip quotes feature

src/main/java/click/kamil/springstatemachineexporter/ast/app/AstTransitionParser.java

@@ -77,7 +77,7 @@ public class AstTransitionParser {
         // Extract source state once
         for (MethodInvocation call : segment) {
             if ("source".equals(call.getName().getIdentifier()) && !call.arguments().isEmpty()) {
-                sourceState = call.arguments().getFirst().toString();
+                sourceState = QuotedExpression.of(call.arguments().getFirst()).toStringWithoutQuotes();
                 break; // only take the first occurrence
             }
         }
@@ -93,11 +93,17 @@ public class AstTransitionParser {
             if (sourceState != null) {
                 t.getSourceStates().add(sourceState);
             }
-            t.getTargetStates().add(args.get(0).toString());
+            t.getTargetStates().add(QuotedExpression.of(args.get(0)).toStringWithoutQuotes());
             if (args.size() > 1) {
-                parseGuard(args.get(1), t);
+                if ("last".equals(methodName)) {
+                    // For 'last', the second argument is an action, not a guard
+                    parseAction(args.get(1), t);
+                } else {
+                    // For 'first' and 'then', the second argument is a guard
+                    parseGuard(args.get(1), t);
+                }
             }
-            if ("then".equals(methodName) && args.size() > 2) {
+            if (args.size() > 2) {
                 parseAction(args.get(2), t);
             }
             t.setOrder(orderCounter++);
@@ -122,9 +128,9 @@ public class AstTransitionParser {
             }
 
             switch (methodName) {
-                case "source" -> args.forEach(arg -> t.getSourceStates().add(arg.toString()));
+                case "source" -> args.forEach(arg -> t.getSourceStates().add(QuotedExpression.of(arg).toStringWithoutQuotes()));
-                case "target" -> args.forEach(arg -> t.getTargetStates().add(arg.toString()));
+                case "target" -> args.forEach(arg -> t.getTargetStates().add(QuotedExpression.of(arg).toStringWithoutQuotes()));
-                case "event" -> t.setEvent(args.getFirst().toString());
+                case "event" -> t.setEvent(QuotedExpression.of(args.getFirst()).toStringWithoutQuotes());
                 case "guard" -> parseGuard(args.getFirst(), t);
                 case "action" -> parseAction(args.getFirst(), t);
             }
@@ -133,16 +139,18 @@ public class AstTransitionParser {
     }
 
     private static void parseGuard(Object arg, Transition t) {
-        if (arg instanceof Expression expr) {
+        QuotedExpression quotedExpr = QuotedExpression.of(arg);
-            t.setGuard(expr.toString());
+        if (quotedExpr != null) {
-            t.setLambdaGuard(isLambdaOrAnonymous(expr));
+            t.setGuard(quotedExpr.toStringWithoutQuotes());
+            t.setLambdaGuard(isLambdaOrAnonymous(quotedExpr.getExpression()));
         }
     }
 
     private static void parseAction(Object arg, Transition t) {
-        if (arg instanceof Expression expr) {
+        QuotedExpression quotedExpr = QuotedExpression.of(arg);
-            t.getActions().add(expr.toString());
+        if (quotedExpr != null) {
-            t.getIsLambdaActions().add(isLambdaOrAnonymous(expr));
+            t.getActions().add(quotedExpr.toStringWithoutQuotes());
+            t.getIsLambdaActions().add(isLambdaOrAnonymous(quotedExpr.getExpression()));
         }
     }
 

src/main/java/click/kamil/springstatemachineexporter/ast/app/QuotedExpression.java

@@ -0,0 +1,34 @@
+package click.kamil.springstatemachineexporter.ast.app;
+
+import lombok.Getter;
+import org.eclipse.jdt.core.dom.Expression;
+
+@Getter
+public class QuotedExpression {
+    private final Expression expression;
+    
+    private QuotedExpression(Expression expression) {
+        this.expression = expression;
+    }
+
+    public String toStringWithoutQuotes() {
+        if (expression == null) return null;
+        return stripQuotes(expression.toString());
+    }
+
+    public static String stripQuotes(String s) {
+        if (s == null) return null;
+        return s.replaceAll("^\"|\"$", "");
+    }
+    
+    public static QuotedExpression of(Expression expression) {
+        return new QuotedExpression(expression);
+    }
+    
+    public static QuotedExpression of(Object obj) {
+        if (obj instanceof Expression expr) {
+            return new QuotedExpression(expr);
+        }
+        return null;
+    }
+} 

src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineStateConfigurationMethodFinder.java

@@ -94,10 +94,10 @@ public class StateMachineStateConfigurationMethodFinder {
                 switch (methodName) {
                     case "initial" -> {
                         if (!isInsideRegionOrFork(call)) {
-                            initialStates.add(arg.toString());
+                            initialStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
                         }
                     }
-                    case "end" -> endStates.add(arg.toString());
+                    case "end" -> endStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
                     // ignore others
                 }
             }

src/main/java/click/kamil/springstatemachineexporter/ast/app/TransitionStateUtils.java

@@ -31,12 +31,7 @@ public class TransitionStateUtils {
 
     private static Stream<String> normalizeStates(Collection<String> states) {
         return states.stream()
-                .map(TransitionStateUtils::stripQuotes)
+                .map(QuotedExpression::stripQuotes)
                 .filter(Objects::nonNull);
     }
-
-    private static String stripQuotes(String s) {
-        if (s == null) return null;
-        return s.replaceAll("^\"|\"$", "");
-    }
 }

src/test/java/click/kamil/springstatemachineexporter/ast/app/AstTransitionParserTest.java

@@ -0,0 +1,438 @@
+package click.kamil.springstatemachineexporter.ast.app;
+
+import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import org.eclipse.jdt.core.dom.*;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class AstTransitionParserTest {
+
+    @Test
+    void shouldParseSimpleWithExternalTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getType()).isEqualTo("withExternal");
+        assertThat(transition.getSourceStates()).containsExactly("CREATED");
+        assertThat(transition.getTargetStates()).containsExactly("PAID");
+        assertThat(transition.getEvent()).isEqualTo("PAY");
+    }
+
+    @Test
+    void shouldParseTransitionWithGuard() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").guard("amount > 0");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getGuard()).isEqualTo("amount > 0");
+        assertThat(transition.isLambdaGuard()).isFalse();
+    }
+
+    @Test
+    void shouldParseTransitionWithLambdaGuard() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").guard(context -> context.getExtendedState().getVariables().get("amount") > 0);
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getGuard()).contains("context -> context.getExtendedState().getVariables().get(\"amount\") > 0");
+        assertThat(transition.isLambdaGuard()).isTrue();
+    }
+
+    @Test
+    void shouldParseTransitionWithAction() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").action("processPayment()");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getActions()).containsExactly("processPayment()");
+        assertThat(transition.getIsLambdaActions()).containsExactly(false);
+    }
+
+    @Test
+    void shouldParseTransitionWithLambdaAction() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").action(context -> System.out.println("Payment processed"));
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getActions()).hasSize(1);
+        assertThat(transition.getActions().getFirst()).contains("context -> System.out.println(\"Payment processed\")");
+        assertThat(transition.getIsLambdaActions()).containsExactly(true);
+    }
+
+    @Test
+    void shouldParseMultipleTransitionsWithAnd() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").and().withExternal().source("PAID").target("SHIPPED").event("SHIP");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(2);
+        
+        Transition first = transitions.getFirst();
+        assertThat(first.getType()).isEqualTo("withExternal");
+        assertThat(first.getSourceStates()).containsExactly("CREATED");
+        assertThat(first.getTargetStates()).containsExactly("PAID");
+        assertThat(first.getEvent()).isEqualTo("PAY");
+        
+        Transition second = transitions.get(1);
+        assertThat(second.getType()).isEqualTo("withExternal");
+        assertThat(second.getSourceStates()).containsExactly("PAID");
+        assertThat(second.getTargetStates()).containsExactly("SHIPPED");
+        assertThat(second.getEvent()).isEqualTo("SHIP");
+    }
+
+    @Test
+    void shouldParseWithChoiceTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withChoice().source("PAID").first("SHIPPED", "isExpress()").then("CANCELLED", "isCancelled()").last("PENDING");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(3);
+        
+        Transition first = transitions.getFirst();
+        assertThat(first.getType()).isEqualTo("withChoice");
+        assertThat(first.getSourceStates()).containsExactly("PAID");
+        assertThat(first.getTargetStates()).containsExactly("SHIPPED");
+        assertThat(first.getGuard()).isEqualTo("isExpress()");
+        assertThat(first.getOrder()).isEqualTo(0);
+        
+        Transition second = transitions.get(1);
+        assertThat(second.getType()).isEqualTo("withChoice");
+        assertThat(second.getSourceStates()).containsExactly("PAID");
+        assertThat(second.getTargetStates()).containsExactly("CANCELLED");
+        assertThat(second.getGuard()).isEqualTo("isCancelled()");
+        assertThat(second.getOrder()).isEqualTo(1);
+        
+        Transition third = transitions.get(2);
+        assertThat(third.getType()).isEqualTo("withChoice");
+        assertThat(third.getSourceStates()).containsExactly("PAID");
+        assertThat(third.getTargetStates()).containsExactly("PENDING");
+        assertThat(third.getOrder()).isEqualTo(2);
+    }
+
+    @Test
+    void shouldParseWithJunctionTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withJunction().source("PAID").first("SHIPPED", "isExpress()").then("CANCELLED", "isCancelled()").last("PENDING");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(3);
+        assertThat(transitions.getFirst().getType()).isEqualTo("withJunction");
+        assertThat(transitions.get(1).getType()).isEqualTo("withJunction");
+        assertThat(transitions.get(2).getType()).isEqualTo("withJunction");
+    }
+
+    @Test
+    void shouldParseWithInternalTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withInternal().source("PAID").event("UPDATE").action("updateStatus()");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getType()).isEqualTo("withInternal");
+        assertThat(transition.getSourceStates()).containsExactly("PAID");
+        assertThat(transition.getEvent()).isEqualTo("UPDATE");
+        assertThat(transition.getActions()).containsExactly("updateStatus()");
+    }
+
+    @Test
+    void shouldParseWithLocalTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withLocal().source("PAID").target("PROCESSING").event("PROCESS");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getType()).isEqualTo("withLocal");
+        assertThat(transition.getSourceStates()).containsExactly("PAID");
+        assertThat(transition.getTargetStates()).containsExactly("PROCESSING");
+        assertThat(transition.getEvent()).isEqualTo("PROCESS");
+    }
+
+    @Test
+    void shouldParseWithForkTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withFork().source("PAID").target("SHIPPING").target("BILLING");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getType()).isEqualTo("withFork");
+        assertThat(transition.getSourceStates()).containsExactly("PAID");
+        assertThat(transition.getTargetStates()).containsExactlyInAnyOrder("SHIPPING", "BILLING");
+    }
+
+    @Test
+    void shouldParseWithJoinTransition() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withJoin().source("SHIPPING").source("BILLING").target("COMPLETED");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getType()).isEqualTo("withJoin");
+        assertThat(transition.getSourceStates()).containsExactlyInAnyOrder("SHIPPING", "BILLING");
+        assertThat(transition.getTargetStates()).containsExactly("COMPLETED");
+    }
+
+    @Test
+    void shouldReturnEmptyList_whenMethodIsNull() {
+        List<Transition> transitions = AstTransitionParser.parseTransitions(null);
+
+        assertThat(transitions).isEmpty();
+    }
+
+    @Test
+    void shouldReturnEmptyList_whenMethodHasNoBody() {
+        String source = """
+                public class TestClass {
+                    public void configure() {
+                        // Empty method
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).isEmpty();
+    }
+
+    @Test
+    void shouldReturnEmptyList_whenMethodHasNoTransitionStatements() {
+        String source = """
+                public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                    System.out.println("No transitions here");
+                    int x = 5;
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).isEmpty();
+    }
+
+    @Test
+    void shouldParseTransitionWithMultipleSources() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").source("PENDING").target("PAID").event("PAY");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getSourceStates()).containsExactlyInAnyOrder("CREATED", "PENDING");
+        assertThat(transition.getTargetStates()).containsExactly("PAID");
+    }
+
+    @Test
+    void shouldParseTransitionWithMultipleTargets() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("PAID").target("SHIPPED").target("CANCELLED").event("PROCESS");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getSourceStates()).containsExactly("PAID");
+        assertThat(transition.getTargetStates()).containsExactlyInAnyOrder("SHIPPED", "CANCELLED");
+    }
+
+    @Test
+    void shouldParseChoiceTransitionWithAction() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions
+                            .withChoice()
+                            .source("PAID")
+                            .first("SHIPPED", "isExpress()", "shipExpress()")
+                            .then("CANCELLED", "isCancelled()", "cancelOrder()")
+                            .last("PENDING", "defaultAction()");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(3);
+        assertThat(transitions.getFirst().getActions()).containsExactly("shipExpress()");
+        assertThat(transitions.get(1).getActions()).containsExactly("cancelOrder()");
+        assertThat(transitions.get(2).getActions()).containsExactly("defaultAction()");
+    }
+
+    @Test
+    void shouldDetectLambdaExpressions() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").guard(context -> context.getExtendedState().getVariables().get("amount") > 0).action(context -> System.out.println("Payment processed"));
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.isLambdaGuard()).isTrue();
+        assertThat(transition.getIsLambdaActions()).containsExactly(true);
+    }
+
+    @Test
+    void shouldDetectAnonymousClassExpressions() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions
+                            .withExternal()
+                            .source("CREATED")
+                            .target("PAID")
+                            .event("PAY")
+                            .action(new Action<String, String>() {
+                                @Override
+                                public void execute(StateContext<String, String> context) {
+                                    System.out.println("Payment processed");
+                                }
+                            });
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+
+        assertThat(transitions).hasSize(1);
+        Transition transition = transitions.getFirst();
+        assertThat(transition.getIsLambdaActions()).containsExactly(true);
+    }
+
+    private MethodDeclaration createMethodDeclaration(String source) {
+        ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
+        parser.setKind(ASTParser.K_COMPILATION_UNIT);
+        parser.setSource(source.toCharArray());
+        CompilationUnit cu = (CompilationUnit) parser.createAST(null);
+        if (cu.types().isEmpty()) {
+            return null;
+        }
+        TypeDeclaration typeDecl = (TypeDeclaration) cu.types().get(0);
+        MethodDeclaration[] methods = typeDecl.getMethods();
+        if (methods.length > 0) {
+            return methods[0];
+        }
+        return null;
+    }
+} 

src/test/java/click/kamil/springstatemachineexporter/ast/app/StateMachineStateConfigurationMethodFinderTest.java

@@ -0,0 +1,392 @@
+package click.kamil.springstatemachineexporter.ast.app;
+
+import org.eclipse.jdt.core.dom.CompilationUnit;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class StateMachineStateConfigurationMethodFinderTest {
+
+    @Test
+    void shouldParseInitialStates() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .end(OrderStates.COMPLETED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("OrderStates.CREATED");
+        assertThat(visitor.getEndStates()).containsExactly("OrderStates.COMPLETED");
+    }
+
+    @Test
+    void shouldParseMultipleInitialStates() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .initial(OrderStates.PENDING)
+                            .end(OrderStates.COMPLETED)
+                            .end(OrderStates.CANCELLED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactlyInAnyOrder("OrderStates.CREATED", "OrderStates.PENDING");
+        assertThat(visitor.getEndStates()).containsExactlyInAnyOrder("OrderStates.COMPLETED", "OrderStates.CANCELLED");
+    }
+
+    @Test
+    void shouldIgnoreStatesInsideFork() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .fork(OrderStates.PROCESSING)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.PROCESSING)
+                            .initial(OrderStates.SHIPPING)
+                            .initial(OrderStates.BILLING)
+                            .end(OrderStates.COMPLETED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("OrderStates.CREATED");
+        assertThat(visitor.getEndStates()).containsExactly("OrderStates.COMPLETED");
+    }
+
+    @Test
+    void shouldIgnoreStatesInsideRegion() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .region(OrderStates.PROCESSING)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.PROCESSING)
+                            .initial(OrderStates.SHIPPING)
+                            .initial(OrderStates.BILLING)
+                            .end(OrderStates.COMPLETED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("OrderStates.CREATED");
+        assertThat(visitor.getEndStates()).containsExactly("OrderStates.COMPLETED");
+    }
+
+    @Test
+    void shouldParseStatesWithStringLiterals() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<String, String> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<String, String> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial("CREATED")
+                            .end("COMPLETED");
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("CREATED");
+        assertThat(visitor.getEndStates()).containsExactly("COMPLETED");
+    }
+
+    @Test
+    void shouldParseStatesWithQuotedStrings() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<String, String> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<String, String> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial("CREATED")
+                            .end("COMPLETED");
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("CREATED");
+        assertThat(visitor.getEndStates()).containsExactly("COMPLETED");
+    }
+
+    @Test
+    void shouldIgnoreNonConfigureMethods() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    public void someOtherMethod() {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED);
+                    }
+                    
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.PENDING)
+                            .end(OrderStates.COMPLETED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("OrderStates.PENDING");
+        assertThat(visitor.getEndStates()).containsExactly("OrderStates.COMPLETED");
+    }
+
+    @Test
+    void shouldHandleEmptyConfigureMethod() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        // Empty method
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).isEmpty();
+        assertThat(visitor.getEndStates()).isEmpty();
+    }
+
+    @Test
+    void shouldHandleConfigureMethodWithoutWithStates() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        System.out.println("No state configuration here");
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).isEmpty();
+        assertThat(visitor.getEndStates()).isEmpty();
+    }
+
+    @Test
+    void shouldParseComplexStateConfiguration() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .state(OrderStates.PAID)
+                            .state(OrderStates.SHIPPED)
+                            .end(OrderStates.COMPLETED)
+                            .end(OrderStates.CANCELLED)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.PAID)
+                            .initial(OrderStates.PROCESSING)
+                            .end(OrderStates.READY);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactlyInAnyOrder("OrderStates.CREATED", "OrderStates.PROCESSING");
+        assertThat(visitor.getEndStates()).containsExactlyInAnyOrder("OrderStates.COMPLETED", "OrderStates.CANCELLED", "OrderStates.READY");
+    }
+
+    @Test
+    void shouldParseStatesWithMethodCalls() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(getInitialState())
+                            .end(getEndState());
+                    }
+                    
+                    private OrderStates getInitialState() {
+                        return OrderStates.CREATED;
+                    }
+                    
+                    private OrderStates getEndState() {
+                        return OrderStates.COMPLETED;
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("getInitialState()");
+        assertThat(visitor.getEndStates()).containsExactly("getEndState()");
+    }
+
+    @Test
+    void shouldHandleMultipleConfigureMethods() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .end(OrderStates.COMPLETED);
+                    }
+                    
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        // This should be ignored
+                        states
+                            .withStates()
+                            .initial(OrderStates.PENDING);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactlyInAnyOrder("OrderStates.CREATED", "OrderStates.PENDING");
+        assertThat(visitor.getEndStates()).containsExactly("OrderStates.COMPLETED");
+    }
+
+    @Test
+    void shouldParseStatesWithEnumValues() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .state(OrderStates.PAID)
+                            .state(OrderStates.SHIPPED)
+                            .end(OrderStates.COMPLETED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("OrderStates.CREATED");
+        assertThat(visitor.getEndStates()).containsExactly("OrderStates.COMPLETED");
+    }
+
+    @Test
+    void shouldHandleNestedForkAndRegion() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        states
+                            .withStates()
+                            .initial(OrderStates.CREATED)
+                            .fork(OrderStates.PROCESSING)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.PROCESSING)
+                            .region(OrderStates.SHIPPING)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.SHIPPING)
+                            .initial(OrderStates.PREPARING)
+                            .end(OrderStates.SHIPPED)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.PROCESSING)
+                            .region(OrderStates.BILLING)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.BILLING)
+                            .initial(OrderStates.CALCULATING)
+                            .end(OrderStates.BILLED)
+                            .and()
+                            .withStates()
+                            .parent(OrderStates.PROCESSING)
+                            .end(OrderStates.COMPLETED);
+                    }
+                }
+                """;
+
+        CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+        StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = 
+            new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+        cu.accept(visitor);
+
+        assertThat(visitor.getInitialStates()).containsExactly("OrderStates.CREATED");
+        assertThat(visitor.getEndStates()).containsExactlyInAnyOrder("OrderStates.COMPLETED", "OrderStates.SHIPPED", "OrderStates.BILLED");
+    }
+} 

src/test/java/click/kamil/springstatemachineexporter/ast/app/StateMachineTransitionConfigurationMethodFinderTest.java

@@ -0,0 +1,344 @@
+package click.kamil.springstatemachineexporter.ast.app;
+
+import org.eclipse.jdt.core.dom.MethodDeclaration;
+import org.eclipse.jdt.core.dom.SimpleName;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class StateMachineTransitionConfigurationMethodFinderTest {
+
+    @Test
+    void shouldFindConfigureMethodInEnumStateMachineConfigurerAdapter() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        transitions
+                            .withExternal()
+                            .source(OrderStates.CREATED)
+                            .target(OrderStates.PAID)
+                            .event(OrderEvents.PAY);
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldFindConfigureMethodInStateMachineConfigurerAdapter() {
+        String source = """
+                public class OrderStateMachineConfig extends StateMachineConfigurerAdapter<String, String> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        transitions
+                            .withExternal()
+                            .source("CREATED")
+                            .target("PAID")
+                            .event("PAY");
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldReturnNull_whenClassDoesNotExtendStateMachineConfigurer() {
+        String source = """
+                public class OrderStateMachineConfig {
+                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
+                        // This should not be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenClassIsInterface() {
+        String source = """
+                public interface OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception;
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenConfigureMethodIsNotPublic() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    protected void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        // This should not be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenConfigureMethodDoesNotReturnVoid() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public String configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        // This should not be found
+                        return "result";
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenConfigureMethodHasWrongNumberOfParameters() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure() throws Exception {
+                        // This should not be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenConfigureMethodParameterIsNotStateMachineTransitionConfigurer() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        // This should not be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenConfigureMethodDoesNotThrowException() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) {
+                        // This should not be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldReturnNull_whenConfigureMethodThrowsDifferentException() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws RuntimeException {
+                        // This should not be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldFindConfigureMethodWithMultipleExceptions() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception, RuntimeException {
+                        transitions
+                            .withExternal()
+                            .source(OrderStates.CREATED)
+                            .target(OrderStates.PAID)
+                            .event(OrderEvents.PAY);
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldFindConfigureMethodWithParameterizedType() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        transitions
+                            .withExternal()
+                            .source(OrderStates.CREATED)
+                            .target(OrderStates.PAID)
+                            .event(OrderEvents.PAY);
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldReturnNull_whenMultipleConfigureMethodsExist() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    public void configure(StateMachineStateConfigurer<OrderStates, OrderEvents> states) throws Exception {
+                        // State configuration
+                    }
+                
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        // Transition configuration
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldHandleClassWithNoMethods() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    // No methods
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method).isNull();
+    }
+
+    @Test
+    void shouldHandleClassWithOtherMethods() {
+        String source = """
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    public void someOtherMethod() {
+                        // This should be ignored
+                    }
+                
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        // This should be found
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldHandleClassWithGenericParameters() {
+        String source = """
+                public class OrderStateMachineConfig<S, E> extends EnumStateMachineConfigurerAdapter<S, E> {
+                    @Override
+                    public void configure(StateMachineTransitionConfigurer<S, E> transitions) throws Exception {
+                        // Generic configuration
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+
+    @Test
+    void shouldHandleClassWithAnnotations() {
+        String source = """
+                @Configuration
+                public class OrderStateMachineConfig extends EnumStateMachineConfigurerAdapter<OrderStates, OrderEvents> {
+                    @Override
+                    @Bean
+                    public void configure(StateMachineTransitionConfigurer<OrderStates, OrderEvents> transitions) throws Exception {
+                        // Configuration with annotations
+                    }
+                }
+                """;
+
+        StateMachineTransitionConfigurationMethodFinder finder = new StateMachineTransitionConfigurationMethodFinder(source);
+        MethodDeclaration method = finder.findConfigureMethod();
+
+        assertThat(method)
+                .extracting(MethodDeclaration::getName)
+                .extracting(SimpleName::getIdentifier)
+                .isEqualTo("configure");
+    }
+} 

src/test/java/click/kamil/springstatemachineexporter/ast/app/TransitionStateUtilsTest.java

@@ -0,0 +1,292 @@
+package click.kamil.springstatemachineexporter.ast.app;
+
+import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Set;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class TransitionStateUtilsTest {
+
+    @Test
+    void shouldFindStartStates_whenStatesHaveNoIncomingTransitions() {
+        List<Transition> transitions = List.of(
+                createTransition("CREATED", "PAID"),
+                createTransition("PAID", "SHIPPED"),
+                createTransition("SHIPPED", "COMPLETED")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+
+        assertThat(startStates).containsExactly("CREATED");
+    }
+
+    @Test
+    void shouldFindStartStates_whenMultipleStatesHaveNoIncomingTransitions() {
+        List<Transition> transitions = List.of(
+                createTransition("CREATED", "PAID"),
+                createTransition("PENDING", "PAID"),
+                createTransition("PAID", "SHIPPED")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+
+        assertThat(startStates).containsExactlyInAnyOrder("CREATED", "PENDING");
+    }
+
+    @Test
+    void shouldFindEndStates_whenStatesHaveNoOutgoingTransitions() {
+        List<Transition> transitions = List.of(
+                createTransition("CREATED", "PAID"),
+                createTransition("PAID", "SHIPPED"),
+                createTransition("PAID", "CANCELLED")
+        );
+
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(endStates).containsExactlyInAnyOrder("SHIPPED", "CANCELLED");
+    }
+
+    @Test
+    void shouldFindEndStates_whenMultipleStatesHaveNoOutgoingTransitions() {
+        List<Transition> transitions = List.of(
+                createTransition("CREATED", "PAID"),
+                createTransition("PAID", "SHIPPED"),
+                createTransition("PAID", "CANCELLED"),
+                createTransition("PAID", "REFUNDED")
+        );
+
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(endStates).containsExactlyInAnyOrder("SHIPPED", "CANCELLED", "REFUNDED");
+    }
+
+    @Test
+    void shouldHandleEmptyTransitionsList() {
+        Set<String> startStates = TransitionStateUtils.findStartStates(List.of());
+        Set<String> endStates = TransitionStateUtils.findEndStates(List.of());
+
+        assertThat(startStates).isEmpty();
+        assertThat(endStates).isEmpty();
+    }
+
+    @Test
+    void shouldThrowException_whenNullTransitionsList() {
+        List<Transition> transitions = null;
+        assertThatThrownBy(() -> TransitionStateUtils.findStartStates(transitions))
+                .isInstanceOf(NullPointerException.class);
+        assertThatThrownBy(() -> TransitionStateUtils.findEndStates(transitions))
+                .isInstanceOf(NullPointerException.class);
+    }
+
+    @Test
+    void shouldHandleTransitionsWithQuotedStates() {
+        List<Transition> transitions = List.of(
+                createTransition("\"CREATED\"", "\"PAID\""),
+                createTransition("\"PAID\"", "\"SHIPPED\"")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("CREATED");
+        assertThat(endStates).containsExactly("SHIPPED");
+    }
+
+    @Test
+    void shouldHandleTransitionsWithMixedQuotedAndUnquotedStates() {
+        List<Transition> transitions = List.of(
+                createTransition("CREATED", "\"PAID\""),
+                createTransition("\"PAID\"", "SHIPPED")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("CREATED");
+        assertThat(endStates).containsExactly("SHIPPED");
+    }
+
+    @Test
+    void shouldHandleTransitionsWithNullStates() {
+        List<Transition> transitions = List.of(
+                createTransition(null, "PAID"),
+                createTransition("PAID", null)
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).isEmpty();
+        assertThat(endStates).isEmpty();
+    }
+
+    @Test
+    void shouldHandleTransitionsWithEmptyStates() {
+        List<Transition> transitions = List.of(
+                createTransition("", "PAID"),
+                createTransition("PAID", "")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).isEmpty();
+        assertThat(endStates).isEmpty();
+    }
+
+    @Test
+    void shouldHandleTransitionsWithMultipleSourceStates() {
+        Transition transition1 = createTransition(List.of("CREATED", "PENDING"), List.of("PAID"));
+        Transition transition2 = createTransition(List.of("PAID"), List.of("SHIPPED"));
+        List<Transition> transitions = List.of(transition1, transition2);
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactlyInAnyOrder("CREATED", "PENDING");
+        assertThat(endStates).containsExactly("SHIPPED");
+    }
+
+    @Test
+    void shouldHandleTransitionsWithMultipleTargetStates() {
+        Transition transition1 = createTransition(
+                List.of("PAID"),
+                List.of("SHIPPED", "CANCELLED")
+        );
+        List<Transition> transitions = List.of(transition1);
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("PAID");
+        assertThat(endStates).containsExactlyInAnyOrder("SHIPPED", "CANCELLED");
+    }
+
+    @Test
+    void shouldHandleCircularTransitions() {
+        List<Transition> transitions = List.of(
+                createTransition("A", "B"),
+                createTransition("B", "C"),
+                createTransition("C", "A")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).isEmpty();
+        assertThat(endStates).isEmpty();
+    }
+
+    @Test
+    void shouldHandleSelfLoopingTransitions() {
+        List<Transition> transitions = List.of(
+                createTransition("A", "A"),
+                createTransition("A", "B")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).isEmpty(); // A is both source and target, so not a start state
+        assertThat(endStates).containsExactly("B"); // Only B is a pure end state
+    }
+
+    @Test
+    void shouldHandleTransitionsWithWhitespaceInStates() {
+        List<Transition> transitions = List.of(
+                createTransition(" CREATED ", " PAID "),
+                createTransition(" PAID ", " SHIPPED ")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly(" CREATED ");
+        assertThat(endStates).containsExactly(" SHIPPED ");
+    }
+
+    @Test
+    void shouldHandleTransitionsWithSpecialCharactersInStates() {
+        List<Transition> transitions = List.of(
+                createTransition("STATE_1", "STATE_2"),
+                createTransition("STATE_2", "STATE_3")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("STATE_1");
+        assertThat(endStates).containsExactly("STATE_3");
+    }
+
+    @Test
+    void shouldHandleTransitionsWithEnumValues() {
+        List<Transition> transitions = List.of(
+                createTransition("OrderStates.CREATED", "OrderStates.PAID"),
+                createTransition("OrderStates.PAID", "OrderStates.SHIPPED")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("OrderStates.CREATED");
+        assertThat(endStates).containsExactly("OrderStates.SHIPPED");
+    }
+
+    @Test
+    void shouldHandleTransitionsWithMethodCalls() {
+        List<Transition> transitions = List.of(
+                createTransition("getInitialState()", "getPaidState()"),
+                createTransition("getPaidState()", "getShippedState()")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("getInitialState()");
+        assertThat(endStates).containsExactly("getShippedState()");
+    }
+
+    @Test
+    void shouldHandleComplexStateMachine() {
+        List<Transition> transitions = List.of(
+                createTransition("CREATED", "PAID"),
+                createTransition("PAID", "SHIPPED"),
+                createTransition("PAID", "CANCELLED"),
+                createTransition("SHIPPED", "DELIVERED"),
+                createTransition("CANCELLED", "REFUNDED")
+        );
+
+        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
+        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
+
+        assertThat(startStates).containsExactly("CREATED");
+        assertThat(endStates).containsExactlyInAnyOrder("DELIVERED", "REFUNDED");
+    }
+
+    private Transition createTransition(String source, String target) {
+        Transition transition = new Transition();
+        if (source != null) {
+            transition.getSourceStates().add(source);
+        }
+        if (target != null) {
+            transition.getTargetStates().add(target);
+        }
+        return transition;
+    }
+
+    private Transition createTransition(List<String> sources, List<String> targets) {
+        Transition transition = new Transition();
+        if (sources != null) {
+            transition.getSourceStates().addAll(sources);
+        }
+        if (targets != null) {
+            transition.getTargetStates().addAll(targets);
+        }
+        return transition;
+    }
+}