removal and better style

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/Main.java

@@ -2,13 +2,9 @@ package click.kamil.springstatemachineexporter;
 
 import click.kamil.springstatemachineexporter.ast.app.AstTransitionParser;
 import click.kamil.springstatemachineexporter.ast.app.StateMachineAggregator;
-import click.kamil.springstatemachineexporter.ast.app.StateMachineStateConfigurationMethodFinder;
-import click.kamil.springstatemachineexporter.ast.app.StateMachineTransitionConfigurationMethodFinder;
 import click.kamil.springstatemachineexporter.ast.app.TransitionStateUtils;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
-import click.kamil.springstatemachineexporter.ast.common.FileUtils;
-import click.kamil.springstatemachineexporter.ast.in.AstFileFinder;
 import click.kamil.springstatemachineexporter.ast.out.Dot;
 import click.kamil.springstatemachineexporter.ast.out.PlantUml;
 import click.kamil.springstatemachineexporter.ast.out.Scxml;
@@ -27,7 +23,6 @@ import java.util.Set;
 public class Main {
     private static final String START_DIR = ".";
     private static final String OUTPUT_DIR = "./out";
-    private static final String EXTENSION = ".java";
     private static final List<String> TARGET_ANNOTATIONS = List.of("EnableStateMachineFactory", "EnableStateMachine");
     public static final Set<String> STATE_MACHINE_RETURN_TYPES = Set.of("StateMachine", "StateMachineFactory", "StateMachineModelFactory");
 

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

@@ -1,15 +1,19 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.Action;
+import click.kamil.springstatemachineexporter.ast.app.domain.Guard;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 import org.eclipse.jdt.core.dom.*;
 
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class AstTransitionParser {
 
-    private static final Set<String> SUPPORTED_TRANSITION_TYPES = Set.of(
+    private static final Set<String> SUPPORTED_TRANSITION_METHOD_NAMES = Arrays.stream(TransitionType.values())
-            "withExternal", "withInternal", "withLocal", "withFork", "withJoin", "withChoice", "withJunction"
+            .map(TransitionType::getMethodName)
-    );
+            .collect(Collectors.toSet());
 
     public static List<Transition> parseTransitions(MethodDeclaration method) {
         if (method == null) return Collections.emptyList();
@@ -30,7 +34,7 @@ public class AstTransitionParser {
                     return false;
                 }
                 List<MethodInvocation> chain = collectMethodInvocationChain(node);
-                return chain.stream().anyMatch(mi -> SUPPORTED_TRANSITION_TYPES.contains(mi.getName().getIdentifier()));
+                return chain.stream().anyMatch(mi -> SUPPORTED_TRANSITION_METHOD_NAMES.contains(mi.getName().getIdentifier()));
             }
         });
 
@@ -88,18 +92,17 @@ public class AstTransitionParser {
         // Step 3: Parse each segment
         for (int i = 0; i < segments.size(); i++) {
             List<MethodInvocation> segment = segments.get(i);
-            boolean isChoice = segment.stream().anyMatch(mi -> "withChoice".equals(mi.getName().getIdentifier()));
+            
-            boolean isJunction = segment.stream().anyMatch(mi -> "withJunction".equals(mi.getName().getIdentifier()));
+            Optional<TransitionType> segmentType = segment.stream()
+                    .map(mi -> TransitionType.fromMethodName(mi.getName().getIdentifier()))
+                    .flatMap(Optional::stream)
+                    .findFirst();
 
             // If it's a continuation and we are at the first segment, we should be careful
             boolean isFirstSegmentOfContinuation = isContinuation && i == 0;
 
-            if (isChoice || isJunction) {
+            if (segmentType.isPresent() && (segmentType.get() == TransitionType.CHOICE || segmentType.get() == TransitionType.JUNCTION)) {
-                List<Transition> choiceTransitions = parseChoiceOrJunctionTransitions(segment, isChoice ? "withChoice" : "withJunction", argsMap);
+                List<Transition> choiceTransitions = parseChoiceOrJunctionTransitions(segment, segmentType.get(), argsMap);
-                if (isFirstSegmentOfContinuation) {
-                    // For choice continuations, the first transition often has an empty source (inherited from receiver)
-                    // We keep them as standard choice parsing handles it, but we ensure we don't introduce new start states
-                }
                 transitions.addAll(choiceTransitions);
             } else {
                 Transition t = parseStandardTransition(segment, argsMap);
@@ -118,7 +121,7 @@ public class AstTransitionParser {
         return transitions;
     }
 
-    private static List<Transition> parseChoiceOrJunctionTransitions(List<MethodInvocation> segment, String type, Map<String, Expression> argsMap) {
+    private static List<Transition> parseChoiceOrJunctionTransitions(List<MethodInvocation> segment, TransitionType type, Map<String, Expression> argsMap) {
         List<Transition> transitions = new ArrayList<>();
         String sourceState = null;
         int orderCounter = 0;
@@ -174,8 +177,9 @@ public class AstTransitionParser {
             String methodName = call.getName().getIdentifier();
             List<?> args = call.arguments();
 
-            if (SUPPORTED_TRANSITION_TYPES.contains(methodName)) {
+            Optional<TransitionType> type = TransitionType.fromMethodName(methodName);
-                t.setType(methodName);
+            if (type.isPresent()) {
+                t.setType(type.get());
                 continue;
             }
             if (args.isEmpty()) {
@@ -211,16 +215,14 @@ public class AstTransitionParser {
     private static void parseGuard(Object arg, Transition t) {
         QuotedExpression quotedExpr = QuotedExpression.of(arg);
         if (quotedExpr != null) {
-            t.setGuard(quotedExpr.toStringWithoutQuotes());
+            t.setGuard(Guard.of(quotedExpr.toStringWithoutQuotes(), isLambdaOrAnonymous(quotedExpr.getExpression())));
-            t.setLambdaGuard(isLambdaOrAnonymous(quotedExpr.getExpression()));
         }
     }
 
     private static void parseAction(Object arg, Transition t) {
         QuotedExpression quotedExpr = QuotedExpression.of(arg);
         if (quotedExpr != null) {
-            t.getActions().add(quotedExpr.toStringWithoutQuotes());
+            t.getActions().add(Action.of(quotedExpr.toStringWithoutQuotes(), isLambdaOrAnonymous(quotedExpr.getExpression())));
-            t.getIsLambdaActions().add(isLambdaOrAnonymous(quotedExpr.getExpression()));
         }
     }
 

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

@@ -2,6 +2,7 @@ package click.kamil.springstatemachineexporter.ast.app;
 
 import lombok.Getter;
 import org.eclipse.jdt.core.dom.Expression;
+import org.eclipse.jdt.core.dom.StringLiteral;
 
 @Getter
 public class QuotedExpression {
@@ -13,6 +14,9 @@ public class QuotedExpression {
 
     public String toStringWithoutQuotes() {
         if (expression == null) return null;
+        if (expression instanceof StringLiteral sl) {
+            return sl.getLiteralValue();
+        }
         return stripQuotes(expression.toString());
     }
 

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineAggregator.java

@@ -1,15 +1,19 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
 import lombok.Getter;
 import org.eclipse.jdt.core.dom.*;
 
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class StateMachineAggregator {
     private final CodebaseContext context;
 
+    private enum AnalysisGoal { TRANSITIONS, STATES }
+
     public StateMachineAggregator(CodebaseContext context) {
         this.context = context;
     }
@@ -76,18 +80,35 @@ public class StateMachineAggregator {
             } else if (isTransitionChainContinuation(mi, argsMap)) {
                 transitions.addAll(AstTransitionParser.parseTransitionsFromExpression(mi, convertToExpressionMap(argsMap), true));
             } else if (isForEach(mi)) {
-                transitions.addAll(parseForEach(mi, searchStartClass, visitedMethods, argsMap, true, null, null));
+                transitions.addAll(parseForEach(mi, searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
             } else {
                 // Check if it's a call on a lambda parameter or a variable that resolved to a lambda
-                Expression receiver = mi.getExpression();
+                Expression lambdaReceiver = mi.getExpression();
-                if (receiver instanceof SimpleName sn) {
+                if (lambdaReceiver instanceof SimpleName sn) {
                     Object resolved = argsMap.get(sn.getIdentifier());
                     if (resolved instanceof LambdaExpression lambda) {
-                        transitions.addAll(parseLambdaInvocation(lambda, mi.arguments(), searchStartClass, visitedMethods, argsMap, true, null, null));
+                        transitions.addAll(parseLambdaInvocation(lambda, mi.arguments(), searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
                     }
                 }
 
                 MethodDeclaration calledMethod = findMethodInHierarchy(mi.getName().getIdentifier(), searchStartClass);
+                if (calledMethod == null) {
+                    // Try to resolve method in another class if it's a static call or on a known type
+                    Expression miReceiver = mi.getExpression();
+                    if (miReceiver instanceof SimpleName sn) {
+                        TypeDeclaration otherTd = context.getTypeDeclaration(sn.getIdentifier(), (CompilationUnit) searchStartClass.getRoot());
+                        if (otherTd != null) {
+                            calledMethod = findMethodInHierarchy(mi.getName().getIdentifier(), otherTd);
+                            if (calledMethod != null) {
+                                // For calls to other classes, we switch the "searchStartClass" to that class
+                                Map<String, Object> nextArgsMap = buildArgsMap(calledMethod, mi, argsMap);
+                                transitions.addAll(parseTransitionsWithMethodCalls(calledMethod, otherTd, visitedMethods, nextArgsMap));
+                                return transitions; // Handled
+                            }
+                        }
+                    }
+                }
+
                 if (calledMethod != null) {
                     Map<String, Object> nextArgsMap = buildArgsMap(calledMethod, mi, argsMap);
                     transitions.addAll(parseTransitionsWithMethodCalls(calledMethod, searchStartClass, visitedMethods, nextArgsMap));
@@ -101,7 +122,7 @@ public class StateMachineAggregator {
         return mi.getName().getIdentifier().equals("forEach");
     }
 
-    private List<Transition> parseLambdaInvocation(LambdaExpression lambda, List<?> arguments, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, boolean isTransitions, Set<String> initialStates, Set<String> endStates) {
+    private List<Transition> parseLambdaInvocation(LambdaExpression lambda, List<?> arguments, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, AnalysisGoal goal, Set<String> initialStates, Set<String> endStates) {
         Map<String, Object> lambdaArgsMap = new HashMap<>(argsMap);
         List<?> params = lambda.parameters();
         for (int i = 0; i < params.size() && i < arguments.size(); i++) {
@@ -113,7 +134,7 @@ public class StateMachineAggregator {
             }
         }
 
-        if (isTransitions) {
+        if (goal == AnalysisGoal.TRANSITIONS) {
             if (lambda.getBody() instanceof Block block) return parseTransitionsInBlock(block, searchStartClass, visitedMethods, lambdaArgsMap);
             else if (lambda.getBody() instanceof Expression lambdaExpr) return parseTransitionsInExpression(lambdaExpr, searchStartClass, visitedMethods, lambdaArgsMap);
         } else {
@@ -123,7 +144,7 @@ public class StateMachineAggregator {
         return Collections.emptyList();
     }
 
-    private List<Transition> parseForEach(MethodInvocation forEachMi, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, boolean isTransitions, Set<String> initialStates, Set<String> endStates) {
+    private List<Transition> parseForEach(MethodInvocation forEachMi, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, AnalysisGoal goal, Set<String> initialStates, Set<String> endStates) {
         List<Transition> transitions = new ArrayList<>();
         Expression receiver = forEachMi.getExpression();
         List<Expression> elements = unrollCollection(receiver, argsMap);
@@ -136,7 +157,7 @@ public class StateMachineAggregator {
                 Map<String, Object> lambdaArgsMap = new HashMap<>(argsMap);
                 if (!paramName.isEmpty()) lambdaArgsMap.put(paramName, element);
                 
-                if (isTransitions) {
+                if (goal == AnalysisGoal.TRANSITIONS) {
                     if (lambda.getBody() instanceof Block block) transitions.addAll(parseTransitionsInBlock(block, searchStartClass, visitedMethods, lambdaArgsMap));
                     else if (lambda.getBody() instanceof Expression lambdaExpr) transitions.addAll(parseTransitionsInExpression(lambdaExpr, searchStartClass, visitedMethods, lambdaArgsMap));
                 } else {
@@ -148,7 +169,7 @@ public class StateMachineAggregator {
         return transitions;
     }
 
-    private List<Transition> parseEnhancedForLoop(EnhancedForStatement efs, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, boolean isTransitions, Set<String> initialStates, Set<String> endStates) {
+    private List<Transition> parseEnhancedForLoop(EnhancedForStatement efs, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, AnalysisGoal goal, Set<String> initialStates, Set<String> endStates) {
         List<Transition> transitions = new ArrayList<>();
         List<Expression> elements = unrollCollection(efs.getExpression(), argsMap);
         String paramName = efs.getParameter().getName().getIdentifier();
@@ -157,7 +178,7 @@ public class StateMachineAggregator {
             Map<String, Object> loopArgsMap = new HashMap<>(argsMap);
             loopArgsMap.put(paramName, element);
 
-            if (isTransitions) {
+            if (goal == AnalysisGoal.TRANSITIONS) {
                 if (efs.getBody() instanceof Block block) transitions.addAll(parseTransitionsInBlock(block, searchStartClass, visitedMethods, loopArgsMap));
                 else if (efs.getBody() instanceof ExpressionStatement es) transitions.addAll(parseTransitionsInExpression(es.getExpression(), searchStartClass, visitedMethods, loopArgsMap));
             } else {
@@ -217,7 +238,7 @@ public class StateMachineAggregator {
                      }
                  }
              }
-             else if (stmt instanceof EnhancedForStatement efs) transitions.addAll(parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, true, null, null));
+             else if (stmt instanceof EnhancedForStatement efs) transitions.addAll(parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
              else if (stmt instanceof TryStatement ts) transitions.addAll(parseTransitionsInBlock(ts.getBody(), searchStartClass, visitedMethods, argsMap));
         }
         return transitions;
@@ -238,7 +259,7 @@ public class StateMachineAggregator {
                      }
                  }
              }
-             else if (stmt instanceof EnhancedForStatement efs) parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, false, initialStates, endStates);
+             else if (stmt instanceof EnhancedForStatement efs) parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
              else if (stmt instanceof TryStatement ts) parseStatesInBlock(ts.getBody(), searchStartClass, initialStates, endStates, visitedMethods, argsMap);
         }
     }
@@ -246,14 +267,14 @@ public class StateMachineAggregator {
     private void parseStatesInExpression(Expression expr, TypeDeclaration searchStartClass, Set<String> initialStates, Set<String> endStates, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
         if (expr instanceof MethodInvocation mi) {
             if (isWithStatesChain(mi, argsMap)) parseStateChain(mi, initialStates, endStates, argsMap);
-            else if (isForEach(mi)) parseForEach(mi, searchStartClass, visitedMethods, argsMap, false, initialStates, endStates);
+            else if (isForEach(mi)) parseForEach(mi, searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
             else {
                 // Check if it's a call on a lambda parameter or a variable that resolved to a lambda
                 Expression receiver = mi.getExpression();
                 if (receiver instanceof SimpleName sn) {
                     Object resolved = argsMap.get(sn.getIdentifier());
                     if (resolved instanceof LambdaExpression lambda) {
-                        parseLambdaInvocation(lambda, mi.arguments(), searchStartClass, visitedMethods, argsMap, false, initialStates, endStates);
+                        parseLambdaInvocation(lambda, mi.arguments(), searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
                     }
                 }
 
@@ -322,7 +343,7 @@ public class StateMachineAggregator {
 
     private boolean isTransitionChainEntry(MethodInvocation mi) {
         String name = mi.getName().getIdentifier();
-        if (Set.of("withExternal", "withInternal", "withLocal", "withFork", "withJoin", "withChoice", "withJunction").contains(name)) return true;
+        if (TransitionType.fromMethodName(name).isPresent()) return true;
         Expression receiver = mi.getExpression();
         if (receiver instanceof MethodInvocation next) return isTransitionChainEntry(next);
         return false;

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

@@ -1,192 +0,0 @@
-package click.kamil.springstatemachineexporter.ast.app;
-
-import lombok.Getter;
-import org.eclipse.jdt.core.dom.*;
-
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-public class StateMachineStateConfigurationMethodFinder {
-
-    public static CompilationUnit parse(String source) {
-        ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
-        parser.setSource(source.toCharArray());
-        parser.setKind(ASTParser.K_COMPILATION_UNIT);
-        parser.setResolveBindings(false);
-        return (CompilationUnit) parser.createAST(null);
-    }
-
-    @Getter
-    public static class StateConfigurerVisitor extends ASTVisitor {
-
-        private final Set<String> initialStates = new HashSet<>();
-        private final Set<String> endStates = new HashSet<>();
-
-        private boolean visitConfigureMethod(MethodDeclaration node) {
-            if (!"configure".equals(node.getName().getIdentifier())) {
-                return super.visit(node);
-            }
-
-            Block body = node.getBody();
-            if (body == null) return false;
-
-            // Traverse statements inside configure method
-            for (Object stmtObj : body.statements()) {
-                if (!(stmtObj instanceof ExpressionStatement exprStmt)) continue;
-                Expression expr = exprStmt.getExpression();
-                if (expr instanceof MethodInvocation rootCall) {
-
-                    // Check if this chain starts with "withStates"
-                    if (isWithStatesChain(rootCall)) {
-                        parseMethodChain(rootCall);
-                    }
-                }
-            }
-            return false;  // no need to visit further inside configure
-        }
-        private boolean visitBuilderMethod(MethodDeclaration node) {
-            Block body = node.getBody();
-            if (body == null) return false;
-
-            // Traverse statements inside buildStateMachine method
-            for (Object stmtObj : body.statements()) {
-                if (!(stmtObj instanceof ExpressionStatement exprStmt)) continue;
-                Expression expr = exprStmt.getExpression();
-                if (expr instanceof MethodInvocation rootCall) {
-                    // Detect when the chain starts with "configureStates"
-                    // Start looping through the method chain
-                    findConfigureStates(rootCall);
-                }
-            }
-            return false;
-        }
-
-
-        private void findConfigureStates(MethodInvocation methodInvocation) {
-            // Start at the outermost method invocation
-            MethodInvocation currentInvocation = methodInvocation;
-
-            while (currentInvocation != null) {
-                // Check if the current method in the chain is 'configureStates'
-                if ("configureStates".equals(currentInvocation.getName().getIdentifier())) {
-                    // Once found, call the parsing method
-                    parseStateMachineMethodChain(methodInvocation);
-                    return;  // Exit once we've handled the method
-                }
-
-                // Move to the next method invocation in the chain (moving forward)
-                Expression nextExpression = currentInvocation.getExpression();
-
-                // Check if the next expression is another method invocation
-                if (nextExpression instanceof MethodInvocation) {
-                    // Continue to the next method invocation
-                    currentInvocation = (MethodInvocation) nextExpression;
-                } else {
-                    // No more method calls in the chain, break out of the loop
-                    currentInvocation = null;
-                }
-            }
-        }
-
-        // Check for method chain starting with "configureStates"
-        private void parseStateMachineMethodChain(MethodInvocation mi) {
-            List<MethodInvocation> chain = new ArrayList<>();
-            Expression current = mi;
-            while (current instanceof MethodInvocation m) {
-                chain.addFirst(m);
-                current = m.getExpression();
-            }
-
-            for (MethodInvocation call : chain) {
-                String methodName = call.getName().getIdentifier();
-                List<?> args = call.arguments();
-                if (args.isEmpty()) {
-                    continue;
-                }
-                Expression arg = (Expression) args.get(0);
-                switch (methodName) {
-                    case "initial" -> {
-                        if (!isInsideRegionOrFork(call)) {
-                            initialStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
-                        }
-                    }
-                    case "end" -> {
-                        endStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
-                    }
-                    // You can also add cases for other methods if needed (e.g., stateEntry, stateExit)
-                    default -> {}
-                }
-            }
-        }
-
-        @Override
-        public boolean visit(MethodDeclaration node) {
-            // Visit method named "configure" only
-            if ("configure".equals(node.getName().getIdentifier())) {
-                return visitConfigureMethod(node);
-            }
-            // Handle builder method configuration (e.g., for buildStateMachine)
-            if ("buildStateMachine".equals(node.getName().getIdentifier())) {
-                return visitBuilderMethod(node);
-            }
-            return false;
-        }
-
-
-        private boolean isWithStatesChain(MethodInvocation mi) {
-            MethodInvocation current = mi;
-            while (true) {
-                Expression expr = current.getExpression();
-                if ("withStates".equals(current.getName().getIdentifier())) {
-                    return true;
-                }
-                if (expr instanceof MethodInvocation next) {
-                    current = next;
-                } else {
-                    return false;
-                }
-            }
-        }
-
-        private boolean isInsideRegionOrFork(MethodInvocation mi) {
-            Expression expr = mi.getExpression();
-            while (expr instanceof MethodInvocation parent) {
-                String methodName = parent.getName().getIdentifier();
-                if (List.of("fork", "region").contains(methodName)) {
-                    return true;
-                }
-                expr = parent.getExpression();
-            }
-            return false;
-        }
-
-        private void parseMethodChain(MethodInvocation mi) {
-            List<MethodInvocation> chain = new ArrayList<>();
-            Expression current = mi;
-            while (current instanceof MethodInvocation m) {
-                chain.addFirst(m);
-                current = m.getExpression();
-            }
-
-            for (MethodInvocation call : chain) {
-                String methodName = call.getName().getIdentifier();
-                List<?> args = call.arguments();
-                if (args.isEmpty()) {
-                    continue;
-                }
-                Expression arg = (Expression) args.getFirst();
-                switch (methodName) {
-                    case "initial" -> {
-                        if (!isInsideRegionOrFork(call)) {
-                            initialStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
-                        }
-                    }
-                    case "end" -> endStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
-                    // ignore others
-                }
-            }
-        }
-    }
-}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineTransitionConfigurationMethodFinder.java

@@ -1,102 +0,0 @@
-package click.kamil.springstatemachineexporter.ast.app;
-
-import click.kamil.springstatemachineexporter.ast.common.AstUtils;
-import lombok.Getter;
-import org.eclipse.jdt.core.dom.*;
-
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-public class StateMachineTransitionConfigurationMethodFinder {
-    private final CompilationUnit cu;
-    private final ConfigureMethodVisitor visitor = new ConfigureMethodVisitor();
-    private static final List<String> stateMachineConfigurationAdapterSuperclasses = List.of("EnumStateMachineConfigurerAdapter", "StateMachineConfigurerAdapter");
-
-    public StateMachineTransitionConfigurationMethodFinder(String source) {
-        ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
-        parser.setSource(source.toCharArray());
-        parser.setKind(ASTParser.K_COMPILATION_UNIT);
-        parser.setResolveBindings(false);
-
-        this.cu = (CompilationUnit) parser.createAST(null);
-    }
-
-    public MethodDeclaration findConfigureMethod() {
-        this.cu.accept(visitor);
-        return visitor.getConfigureMethod();
-    }
-
-    public Set<MethodDeclaration> findMethodsWithReturnType(Set<String> returnTypes) {
-        Set<MethodDeclaration> matchedMethods = new HashSet<>();
-        cu.accept(new ASTVisitor() {
-            @Override
-            public boolean visit(MethodDeclaration node) {
-                Type returnType = node.getReturnType2();
-                String typeName = (returnType == null) ? "void" : AstUtils.extractSimpleTypeName(returnType);
-                if (returnTypes.contains(typeName)) {
-                    matchedMethods.add(node);
-                }
-                return true;
-            }
-        });
-        return matchedMethods;
-    }
-
-
-    @Getter
-    private static class ConfigureMethodVisitor extends ASTVisitor {
-        private MethodDeclaration configureMethod = null;
-
-        @Override
-        public boolean visit(TypeDeclaration node) {
-            if (!node.isInterface() && extendsStateMachineConfigurerAdapter(node)) {
-                for (MethodDeclaration method : node.getMethods()) {
-                    if (isConfigureMethod(method)) {
-                        configureMethod = method;
-                        return false; // Found, stop visiting further
-                    }
-                }
-            }
-            return true;
-        }
-
-        private boolean extendsStateMachineConfigurerAdapter(TypeDeclaration node) {
-            Type superclass = node.getSuperclassType();
-            if (superclass == null)
-                return false;
-
-            if (superclass.isParameterizedType()) {
-                ParameterizedType pt = (ParameterizedType) superclass;
-                return stateMachineConfigurationAdapterSuperclasses.contains(pt.getType().toString());
-            }
-            return stateMachineConfigurationAdapterSuperclasses.contains(superclass.toString());
-        }
-
-        private boolean isConfigureMethod(MethodDeclaration method) {
-            if (!method.getName().getIdentifier().equals("configure")) return false;
-
-            if (!Modifier.isPublic(method.getModifiers())) return false;
-            Type returnType = method.getReturnType2();
-            if (returnType == null || !returnType.isPrimitiveType() || !"void".equals(returnType.toString()))
-                return false;
-
-            @SuppressWarnings("unchecked")
-            List<SingleVariableDeclaration> params = method.parameters();
-            if (params.size() != 1)
-                return false;
-
-            if (!params.getFirst().getType().toString().startsWith("StateMachineTransitionConfigurer"))
-                return false;
-
-            @SuppressWarnings("unchecked")
-            List<Type> thrownExceptions = method.thrownExceptionTypes();
-            for (Type excType : thrownExceptions) {
-                if (excType.toString().equals("Exception")) {
-                    return true;
-                }
-            }
-            return false;
-        }
-    }
-}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/domain/Action.java

@@ -0,0 +1,7 @@
+package click.kamil.springstatemachineexporter.ast.app.domain;
+
+public record Action(String expression, boolean isLambda) {
+    public static Action of(String expression, boolean isLambda) {
+        return expression != null ? new Action(expression, isLambda) : null;
+    }
+}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/domain/Guard.java

@@ -0,0 +1,7 @@
+package click.kamil.springstatemachineexporter.ast.app.domain;
+
+public record Guard(String expression, boolean isLambda) {
+    public static Guard of(String expression, boolean isLambda) {
+        return expression != null ? new Guard(expression, isLambda) : null;
+    }
+}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/domain/Transition.java

@@ -11,16 +11,13 @@ import java.util.List;
 @Getter
 @Setter
 public class Transition {
-    private String type;         // withExternal, withChoice, withInternal, etc.
+    private TransitionType type;
     private List<String> sourceStates = new ArrayList<>();
     private List<String> targetStates = new ArrayList<>();
     private String event;
 
-    private String guard;
+    private Guard guard;
-    private boolean isLambdaGuard = false;
+    private List<Action> actions = new ArrayList<>();
-
-    private List<String> actions = new ArrayList<>();
-    private List<Boolean> isLambdaActions = new ArrayList<>();
 
     private Integer order = null;  // optional order for withChoice or withJunction
 }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/domain/TransitionType.java

@@ -0,0 +1,27 @@
+package click.kamil.springstatemachineexporter.ast.app.domain;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+
+import java.util.Arrays;
+import java.util.Optional;
+
+@Getter
+@RequiredArgsConstructor
+public enum TransitionType {
+    EXTERNAL("withExternal"),
+    INTERNAL("withInternal"),
+    LOCAL("withLocal"),
+    FORK("withFork"),
+    JOIN("withJoin"),
+    CHOICE("withChoice"),
+    JUNCTION("withJunction");
+
+    private final String methodName;
+
+    public static Optional<TransitionType> fromMethodName(String methodName) {
+        return Arrays.stream(values())
+                .filter(t -> t.methodName.equals(methodName))
+                .findFirst();
+    }
+}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/in/AstFileFinder.java

@@ -1,75 +0,0 @@
-package click.kamil.springstatemachineexporter.ast.in;
-
-import click.kamil.springstatemachineexporter.ast.common.AstUtils;
-import org.eclipse.jdt.core.dom.*;
-
-import java.util.List;
-import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-public class AstFileFinder {
-    private final ASTParser parser;
-
-    public AstFileFinder() {
-        this.parser = ASTParser.newParser(AST.getJLSLatest());
-        parser.setKind(ASTParser.K_COMPILATION_UNIT);
-        parser.setResolveBindings(false);
-    }
-
-    public boolean hasClassWithAnnotation(String source, List<String> annotationNames) {
-        parser.setSource(source.toCharArray());
-        CompilationUnit cu = (CompilationUnit) parser.createAST(null);
-        AtomicBoolean found = new AtomicBoolean(false);
-        cu.accept(new ASTVisitor() {
-            @Override
-            public boolean visit(TypeDeclaration node) {
-                // skip interfaces
-                if (node.isInterface()) return true;
-                for (Object modifier : node.modifiers()) {
-                    if (modifier instanceof Annotation annotation &&
-                            annotationNames.contains(annotation.getTypeName().getFullyQualifiedName())) {
-                        found.set(true);
-                        return false;
-                    }
-                }
-                return true;
-            }
-        });
-        return found.get();
-    }
-    public boolean hasFunctionReturningType(String source, Set<String> targetReturnTypes) {
-        parser.setSource(source.toCharArray());
-        CompilationUnit cu = (CompilationUnit) parser.createAST(null);
-        AtomicBoolean found = new AtomicBoolean(false);
-
-        cu.accept(new ASTVisitor() {
-            @Override
-            public boolean visit(MethodDeclaration node) {
-                Type returnType = node.getReturnType2();
-                String typeName = (returnType == null) ? "void" : AstUtils.extractSimpleTypeName(returnType);
-                
-                if (targetReturnTypes.contains(typeName)) {
-                    // Also check for @Bean annotation
-                    boolean hasBeanAnnotation = false;
-                    for (Object modifier : node.modifiers()) {
-                        if (modifier instanceof Annotation annotation) {
-                            String annotationName = annotation.getTypeName().getFullyQualifiedName();
-                            if ("Bean".equals(annotationName) || "org.springframework.context.annotation.Bean".equals(annotationName)) {
-                                hasBeanAnnotation = true;
-                                break;
-                            }
-                        }
-                    }
-                    if (hasBeanAnnotation) {
-                        found.set(true);
-                        return false; // stop visiting
-                    }
-                }
-
-                return true;
-            }
-        });
-
-        return found.get();
-    }
-}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/Dot.java

@@ -1,6 +1,7 @@
 package click.kamil.springstatemachineexporter.ast.out;
 
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 
 import java.util.*;
 
@@ -29,7 +30,7 @@ public class Dot implements StateMachineExporter {
         int colorIndex = 0;
 
         for (Transition t : transitions) {
-            if ("withChoice".equals(t.getType()) && t.getSourceStates() != null) {
+            if (t.getType() == TransitionType.CHOICE && t.getSourceStates() != null) {
                 for (String source : t.getSourceStates()) {
                     String simplified = simplify(source);
                     statesWithChoice.add(simplified);
@@ -90,7 +91,7 @@ public class Dot implements StateMachineExporter {
         for (Transition t : transitions) {
             if (t.getSourceStates() == null || t.getSourceStates().isEmpty()) continue;
 
-            boolean allowNoTarget = "withJoin".equals(t.getType()) || "withFork".equals(t.getType());
+            boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
             List<String> targets = t.getTargetStates();
 
             if (targets == null || targets.isEmpty()) {
@@ -108,7 +109,7 @@ public class Dot implements StateMachineExporter {
                     String target = simplify(rawTarget);
                     String edgeSource = source;
 
-                    if ("withChoice".equals(t.getType()) && statesWithChoice.contains(source)) {
+                    if (t.getType() == TransitionType.CHOICE && statesWithChoice.contains(source)) {
                         edgeSource = includeChoiceStates ? source + "_choice" : source;
                     }
 
@@ -120,10 +121,10 @@ public class Dot implements StateMachineExporter {
                     }
 
                     // Guard
-                    if (t.getGuard() != null && !t.getGuard().isBlank()) {
+                    if (t.getGuard() != null && !t.getGuard().expression().isBlank()) {
                         if (!label.isEmpty()) label.append(" ");
-                        String guardText = useLambdaGuards && t.isLambdaGuard() ? "λ"
+                        String guardText = useLambdaGuards && t.getGuard().isLambda() ? "λ"
-                                : escapeDotString(t.getGuard().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim());
+                                : escapeDotString(t.getGuard().expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim());
                         label.append("[").append(guardText).append("]");
                     }
 
@@ -132,16 +133,15 @@ public class Dot implements StateMachineExporter {
                         StringBuilder actionBuilder = new StringBuilder();
                         for (int i = 0; i < t.getActions().size(); i++) {
                             if (i > 0) actionBuilder.append(", ");
-                            String action = t.getActions().get(i);
+                            var action = t.getActions().get(i);
-                            boolean isLambda = i < t.getIsLambdaActions().size() && t.getIsLambdaActions().get(i);
+                            actionBuilder.append(useLambdaGuards && action.isLambda() ? "λ" : escapeDotString(action.expression()));
-                            actionBuilder.append(useLambdaGuards && isLambda ? "λ" : escapeDotString(action));
                         }
                         if (!label.isEmpty()) label.append(" / ");
                         label.append(actionBuilder);
                     }
 
                     // Order
-                    if ((t.getType().equals("withChoice") || t.getType().equals("withJunction")) && t.getOrder() != null) {
+                    if ((t.getType() == TransitionType.CHOICE || t.getType() == TransitionType.JUNCTION) && t.getOrder() != null) {
                         if (!label.isEmpty()) label.append(" ");
                         label.append("(order=").append(t.getOrder()).append(")");
                     }
@@ -153,36 +153,41 @@ public class Dot implements StateMachineExporter {
                     }
 
                     // Style/color by type
-                    switch (t.getType()) {
+                    if (t.getType() != null) {
-                        case "withInternal":
+                        switch (t.getType()) {
-                            appendEdgeAttr(edgeAttrs, "style", "dashed");
+                            case INTERNAL:
-                            appendEdgeAttr(edgeAttrs, "color", "gray");
+                                appendEdgeAttr(edgeAttrs, "style", "dashed");
-                            break;
+                                appendEdgeAttr(edgeAttrs, "color", "gray");
-                        case "withLocal":
+                                break;
-                            appendEdgeAttr(edgeAttrs, "style", "dotted");
+                            case LOCAL:
-                            appendEdgeAttr(edgeAttrs, "color", "darkgray");
+                                appendEdgeAttr(edgeAttrs, "style", "dotted");
-                            break;
+                                appendEdgeAttr(edgeAttrs, "color", "darkgray");
-                        case "withChoice":
+                                break;
-                            String choiceColor = choiceColorMap.getOrDefault(source, "blue");
+                            case CHOICE:
-                            appendEdgeAttr(edgeAttrs, "style", "solid");
+                                String choiceColor = choiceColorMap.getOrDefault(source, "blue");
-                            appendEdgeAttr(edgeAttrs, "color", choiceColor);
+                                appendEdgeAttr(edgeAttrs, "style", "solid");
-                            break;
+                                appendEdgeAttr(edgeAttrs, "color", choiceColor);
-                        case "withJunction":
+                                break;
-                            appendEdgeAttr(edgeAttrs, "style", "dotted");
+                            case JUNCTION:
-                            appendEdgeAttr(edgeAttrs, "color", "purple");
+                                appendEdgeAttr(edgeAttrs, "style", "dotted");
-                            break;
+                                appendEdgeAttr(edgeAttrs, "color", "purple");
-                        case "withFork":
+                                break;
-                            appendEdgeAttr(edgeAttrs, "style", "bold");
+                            case FORK:
-                            appendEdgeAttr(edgeAttrs, "color", "green");
+                                appendEdgeAttr(edgeAttrs, "style", "bold");
-                            break;
+                                appendEdgeAttr(edgeAttrs, "color", "green");
-                        case "withJoin":
+                                break;
-                            appendEdgeAttr(edgeAttrs, "style", "bold");
+                            case JOIN:
-                            appendEdgeAttr(edgeAttrs, "color", "darkorange");
+                                appendEdgeAttr(edgeAttrs, "style", "bold");
-                            break;
+                                appendEdgeAttr(edgeAttrs, "color", "darkorange");
-                        case "withExternal":
+                                break;
-                        default:
+                            case EXTERNAL:
-                            appendEdgeAttr(edgeAttrs, "style", "solid");
+                            default:
-                            appendEdgeAttr(edgeAttrs, "color", "black");
+                                appendEdgeAttr(edgeAttrs, "style", "solid");
+                                appendEdgeAttr(edgeAttrs, "color", "black");
+                        }
+                    } else {
+                        appendEdgeAttr(edgeAttrs, "style", "solid");
+                        appendEdgeAttr(edgeAttrs, "color", "black");
                     }
 
                     sb.append("  ").append(edgeSource).append(" -> ").append(target);

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/PlantUml.java

@@ -1,6 +1,7 @@
 package click.kamil.springstatemachineexporter.ast.out;
 
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 import click.kamil.springstatemachineexporter.ast.common.StringUtils;
 
 import java.util.*;
@@ -11,15 +12,15 @@ public class PlantUml implements StateMachineExporter {
     // Color palette for choices
     private final List<String> choiceColors = List.of("FF6347", "4682B4", "32CD32", "FFD700", "6A5ACD", "FF69B4");
     private String getColor(Transition t, String source, Map<String, String> choiceStateColorMap) {
+        if (t.getType() == null) return "000000";
         return switch (t.getType()) {
-            case "withChoice" -> choiceStateColorMap.getOrDefault(source, "000000");
+            case CHOICE -> choiceStateColorMap.getOrDefault(source, "000000");
-            case "withExternal" -> "1E90FF";
+            case EXTERNAL -> "1E90FF";
-            case "withInternal" -> "32CD32";
+            case INTERNAL -> "32CD32";
-            case "withLocal" -> "FFA500";
+            case LOCAL -> "FFA500";
-            case "withJunction" -> "FF69B4";
+            case JUNCTION -> "FF69B4";
-            case "withJoin" -> "8A2BE2";
+            case JOIN -> "8A2BE2";
-            case "withFork" -> "20B2AA";
+            case FORK -> "20B2AA";
-            default -> "000000";
         };
     }
     private final String LAMBDA = "λ";
@@ -45,7 +46,7 @@ public class PlantUml implements StateMachineExporter {
         // 2. Find all states that have choice transitions originating from them
         Set<String> statesWithChoice = new HashSet<>();
         for (Transition t : transitions) {
-            if ("withChoice".equals(t.getType())) {
+            if (t.getType() == TransitionType.CHOICE) {
                 for (String source : t.getSourceStates()) {
                     statesWithChoice.add(simplify(source));
                 }
@@ -79,7 +80,7 @@ public class PlantUml implements StateMachineExporter {
         for (Transition t : transitions) {
             if (t.getSourceStates() == null || t.getSourceStates().isEmpty()) continue;
 
-            boolean allowNoTarget = "withJoin".equals(t.getType()) || "withFork".equals(t.getType());
+            boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
 
             List<String> targets;
             if (t.getTargetStates() == null || t.getTargetStates().isEmpty()) {
@@ -99,7 +100,7 @@ public class PlantUml implements StateMachineExporter {
                     String target = simplify(rawTarget);
                     String transitionSource;
                     if (includeChoiceStates) {
-                        transitionSource = "withChoice".equals(t.getType()) && statesWithChoice.contains(source)
+                        transitionSource = t.getType() == TransitionType.CHOICE && statesWithChoice.contains(source)
                                 ? source + "_choice"
                                 : source;
                     } else  {
@@ -134,14 +135,13 @@ public class PlantUml implements StateMachineExporter {
     }
 
     private String getGuardText(Transition t, boolean useLambdaGuards) {
-        String guard = t.getGuard();
+        if (t.getGuard() == null || StringUtils.isBlank(t.getGuard().expression())) {
-        if (StringUtils.isBlank(guard)) {
             return "";
         }
-        if (useLambdaGuards && t.isLambdaGuard()) {
+        if (useLambdaGuards && t.getGuard().isLambda()) {
             return LAMBDA;
         }
-        return guard
+        return t.getGuard().expression()
                 .replaceAll("[\\n\\r]+", " ")
                 .replaceAll("\\s{2,}", " ")
                 .trim();
@@ -152,14 +152,8 @@ public class PlantUml implements StateMachineExporter {
             return "";
         }
 
-        List<String> actions = t.getActions();
+        return t.getActions().stream()
-        List<Boolean> isLambdaFlags = t.getIsLambdaActions();
+                .map(action -> (useLambdaGuards && action.isLambda()) ? LAMBDA : action.expression())
-
-        // Defensive: ensure lists are same size to avoid IndexOutOfBoundsException
-        int size = Math.min(actions.size(), isLambdaFlags.size());
-
-        return IntStream.range(0, size)
-                .mapToObj(i -> (useLambdaGuards && isLambdaFlags.get(i)) ? LAMBDA : actions.get(i))
                 .collect(Collectors.joining(", "));
     }
 
@@ -185,7 +179,7 @@ public class PlantUml implements StateMachineExporter {
                     }
                 });
         Optional.ofNullable(t.getOrder())
-                .filter(order -> List.of("withChoice", "withJunction").contains(t.getType()))
+                .filter(order -> t.getType() == TransitionType.CHOICE || t.getType() == TransitionType.JUNCTION)
                 .map(order -> "(order=" + order + ")")
                 .ifPresent(parts::add);
 

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/Scxml.java

@@ -1,6 +1,7 @@
 package click.kamil.springstatemachineexporter.ast.out;
 
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 
 import java.util.HashSet;
 import java.util.List;
@@ -40,7 +41,7 @@ public class Scxml implements StateMachineExporter {
             for (Transition t : transitions) {
                 if (t.getSourceStates() == null || !t.getSourceStates().contains(state)) continue;
 
-                boolean allowNoTarget = "withJoin".equals(t.getType()) || "withFork".equals(t.getType());
+                boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
                 List<String> targets;
 
                 if (t.getTargetStates() == null || t.getTargetStates().isEmpty()) {
@@ -74,11 +75,11 @@ public class Scxml implements StateMachineExporter {
     }
 
     private static String getGuardText(Transition t, boolean useLambdaGuards) {
-        if (t.getGuard() == null || t.getGuard().isBlank()) return "";
+        if (t.getGuard() == null || t.getGuard().expression().isBlank()) return "";
-        if (useLambdaGuards && t.isLambdaGuard()) {
+        if (useLambdaGuards && t.getGuard().isLambda()) {
             return LAMBDA;
         }
-        return t.getGuard().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
+        return t.getGuard().expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
     }
 
     private static String escapeXml(String s) {
@@ -102,7 +103,7 @@ public class Scxml implements StateMachineExporter {
         }
 
         // Choice/junction order → XML comment
-        if (("withChoice".equals(t.getType()) || "withJunction".equals(t.getType())) && t.getOrder() != null) {
+        if ((t.getType() == TransitionType.CHOICE || t.getType() == TransitionType.JUNCTION) && t.getOrder() != null) {
             sb.append(">\n");
             sb.append("      <!-- order=").append(t.getOrder()).append(" -->\n");
             sb.append("    </transition>\n");

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

@@ -1,6 +1,9 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.Action;
+import click.kamil.springstatemachineexporter.ast.app.domain.Guard;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
+import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 import org.eclipse.jdt.core.dom.*;
 import org.junit.jupiter.api.Test;
 
@@ -27,7 +30,7 @@ class AstTransitionParserTest {
         Transition transition = transitions.getFirst();
         assertThat(transition)
                 .extracting(Transition::getType)
-                .isEqualTo("withExternal");
+                .isEqualTo(TransitionType.EXTERNAL);
         assertThat(transition.getSourceStates()).containsExactly("CREATED");
         assertThat(transition.getTargetStates()).containsExactly("PAID");
         assertThat(transition)
@@ -50,10 +53,12 @@ class AstTransitionParserTest {
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition)
+        assertThat(transition.getGuard())
-                .extracting(Transition::getGuard)
+                .extracting(Guard::expression)
                 .isEqualTo("amount > 0");
-        assertThat(transition.isLambdaGuard()).isFalse();
+        assertThat(transition.getGuard())
+                .extracting(Guard::isLambda)
+                .isEqualTo(false);
     }
 
     @Test
@@ -71,8 +76,8 @@ class AstTransitionParserTest {
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition.getGuard()).contains("context -> context.getExtendedState().getVariables().get(\"amount\") > 0");
+        assertThat(transition.getGuard().expression()).contains("context -> context.getExtendedState().getVariables().get(\"amount\") > 0");
-        assertThat(transition.isLambdaGuard()).isTrue();
+        assertThat(transition.getGuard().isLambda()).isTrue();
     }
 
     @Test
@@ -90,8 +95,12 @@ class AstTransitionParserTest {
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition.getActions()).containsExactly("processPayment()");
+        assertThat(transition.getActions())
-        assertThat(transition.getIsLambdaActions()).containsExactly(false);
+                .extracting(Action::expression)
+                .containsExactly("processPayment()");
+        assertThat(transition.getActions())
+                .extracting(Action::isLambda)
+                .containsExactly(false);
     }
 
     @Test
@@ -109,9 +118,12 @@ class AstTransitionParserTest {
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition.getActions()).hasSize(1);
+        assertThat(transition.getActions())
-        assertThat(transition.getActions().getFirst()).contains("context -> System.out.println(\"Payment processed\")");
+                .extracting(Action::expression)
-        assertThat(transition.getIsLambdaActions()).containsExactly(true);
+                .anyMatch(s -> s.contains("context -> System.out.println(\"Payment processed\")"));
+        assertThat(transition.getActions())
+                .extracting(Action::isLambda)
+                .containsExactly(true);
     }
 
     @Test
@@ -132,7 +144,7 @@ class AstTransitionParserTest {
         Transition first = transitions.getFirst();
         assertThat(first)
                 .extracting(Transition::getType)
-                .isEqualTo("withExternal");
+                .isEqualTo(TransitionType.EXTERNAL);
         assertThat(first.getSourceStates()).containsExactly("CREATED");
         assertThat(first.getTargetStates()).containsExactly("PAID");
         assertThat(first)
@@ -142,7 +154,7 @@ class AstTransitionParserTest {
         Transition second = transitions.get(1);
         assertThat(second)
                 .extracting(Transition::getType)
-                .isEqualTo("withExternal");
+                .isEqualTo(TransitionType.EXTERNAL);
         assertThat(second.getSourceStates()).containsExactly("PAID");
         assertThat(second.getTargetStates()).containsExactly("SHIPPED");
         assertThat(second)
@@ -168,11 +180,11 @@ class AstTransitionParserTest {
         Transition first = transitions.getFirst();
         assertThat(first)
                 .extracting(Transition::getType)
-                .isEqualTo("withChoice");
+                .isEqualTo(TransitionType.CHOICE);
         assertThat(first.getSourceStates()).containsExactly("PAID");
         assertThat(first.getTargetStates()).containsExactly("SHIPPED");
-        assertThat(first)
+        assertThat(first.getGuard())
-                .extracting(Transition::getGuard)
+                .extracting(Guard::expression)
                 .isEqualTo("isExpress()");
         assertThat(first)
                 .extracting(Transition::getOrder)
@@ -181,11 +193,11 @@ class AstTransitionParserTest {
         Transition second = transitions.get(1);
         assertThat(second)
                 .extracting(Transition::getType)
-                .isEqualTo("withChoice");
+                .isEqualTo(TransitionType.CHOICE);
         assertThat(second.getSourceStates()).containsExactly("PAID");
         assertThat(second.getTargetStates()).containsExactly("CANCELLED");
-        assertThat(second)
+        assertThat(second.getGuard())
-                .extracting(Transition::getGuard)
+                .extracting(Guard::expression)
                 .isEqualTo("isCancelled()");
         assertThat(second)
                 .extracting(Transition::getOrder)
@@ -194,7 +206,7 @@ class AstTransitionParserTest {
         Transition third = transitions.get(2);
         assertThat(third)
                 .extracting(Transition::getType)
-                .isEqualTo("withChoice");
+                .isEqualTo(TransitionType.CHOICE);
         assertThat(third.getSourceStates()).containsExactly("PAID");
         assertThat(third.getTargetStates()).containsExactly("PENDING");
         assertThat(third)
@@ -218,13 +230,13 @@ class AstTransitionParserTest {
         assertThat(transitions).hasSize(3);
         assertThat(transitions.getFirst())
                 .extracting(Transition::getType)
-                .isEqualTo("withJunction");
+                .isEqualTo(TransitionType.JUNCTION);
         assertThat(transitions.get(1))
                 .extracting(Transition::getType)
-                .isEqualTo("withJunction");
+                .isEqualTo(TransitionType.JUNCTION);
         assertThat(transitions.get(2))
                 .extracting(Transition::getType)
-                .isEqualTo("withJunction");
+                .isEqualTo(TransitionType.JUNCTION);
     }
 
     @Test
@@ -244,12 +256,14 @@ class AstTransitionParserTest {
         Transition transition = transitions.getFirst();
         assertThat(transition)
                 .extracting(Transition::getType)
-                .isEqualTo("withInternal");
+                .isEqualTo(TransitionType.INTERNAL);
         assertThat(transition.getSourceStates()).containsExactly("PAID");
         assertThat(transition)
                 .extracting(Transition::getEvent)
                 .isEqualTo("UPDATE");
-        assertThat(transition.getActions()).containsExactly("updateStatus()");
+        assertThat(transition.getActions())
+                .extracting(Action::expression)
+                .containsExactly("updateStatus()");
     }
 
     @Test
@@ -269,7 +283,7 @@ class AstTransitionParserTest {
         Transition transition = transitions.getFirst();
         assertThat(transition)
                 .extracting(Transition::getType)
-                .isEqualTo("withLocal");
+                .isEqualTo(TransitionType.LOCAL);
         assertThat(transition.getSourceStates()).containsExactly("PAID");
         assertThat(transition.getTargetStates()).containsExactly("PROCESSING");
         assertThat(transition)
@@ -294,7 +308,7 @@ class AstTransitionParserTest {
         Transition transition = transitions.getFirst();
         assertThat(transition)
                 .extracting(Transition::getType)
-                .isEqualTo("withFork");
+                .isEqualTo(TransitionType.FORK);
         assertThat(transition.getSourceStates()).containsExactly("PAID");
         assertThat(transition.getTargetStates()).containsExactlyInAnyOrder("SHIPPING", "BILLING");
     }
@@ -316,7 +330,7 @@ class AstTransitionParserTest {
         Transition transition = transitions.getFirst();
         assertThat(transition)
                 .extracting(Transition::getType)
-                .isEqualTo("withJoin");
+                .isEqualTo(TransitionType.JOIN);
         assertThat(transition.getSourceStates()).containsExactlyInAnyOrder("SHIPPING", "BILLING");
         assertThat(transition.getTargetStates()).containsExactly("COMPLETED");
     }
@@ -416,9 +430,15 @@ class AstTransitionParserTest {
         List<Transition> transitions = AstTransitionParser.parseTransitions(method);
 
         assertThat(transitions).hasSize(3);
-        assertThat(transitions.getFirst().getActions()).containsExactly("shipExpress()");
+        assertThat(transitions.getFirst().getActions())
-        assertThat(transitions.get(1).getActions()).containsExactly("cancelOrder()");
+                .extracting(Action::expression)
-        assertThat(transitions.get(2).getActions()).containsExactly("defaultAction()");
+                .containsExactly("shipExpress()");
+        assertThat(transitions.get(1).getActions())
+                .extracting(Action::expression)
+                .containsExactly("cancelOrder()");
+        assertThat(transitions.get(2).getActions())
+                .extracting(Action::expression)
+                .containsExactly("defaultAction()");
     }
 
     @Test
@@ -436,8 +456,10 @@ class AstTransitionParserTest {
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition.isLambdaGuard()).isTrue();
+        assertThat(transition.getGuard().isLambda()).isTrue();
-        assertThat(transition.getIsLambdaActions()).containsExactly(true);
+        assertThat(transition.getActions())
+                .extracting(Action::isLambda)
+                .containsExactly(true);
     }
 
     @Test
@@ -465,7 +487,9 @@ class AstTransitionParserTest {
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition.getIsLambdaActions()).containsExactly(true);
+        assertThat(transition.getActions())
+                .extracting(Action::isLambda)
+                .containsExactly(true);
     }
 
     @Test

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/app/StateMachineAggregatorTest.java

@@ -272,6 +272,172 @@ class StateMachineAggregatorTest {
         assertThat(transitions).anyMatch(t -> t.getEvent().equals("E2"));
     }
 
+    @Test
+    void shouldIgnoreStatesInsideFork() throws IOException {
+        String source = """
+                package com.example;
+                import org.springframework.statemachine.config.builders.StateMachineStateConfigurer;
+                public class ForkStateConfig extends StateMachineConfigurerAdapter {
+                    @Override
+                    public void configure(StateMachineStateConfigurer states) throws Exception {
+                        states
+                            .withStates()
+                            .initial("CREATED")
+                            .fork("PROCESSING")
+                            .and()
+                            .withStates()
+                            .parent("PROCESSING")
+                            .initial("SHIPPING")
+                            .initial("BILLING")
+                            .end("COMPLETED");
+                    }
+                }
+                """;
+
+        writeClass("ForkStateConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("ForkStateConfig");
+        aggregator.aggregateStates(td);
+
+        assertThat(aggregator.getInitialStates()).containsExactly("CREATED");
+        assertThat(aggregator.getEndStates()).containsExactly("COMPLETED");
+    }
+
+    @Test
+    void shouldHandleNestedForkAndRegion() throws IOException {
+        String source = """
+                package com.example;
+                import org.springframework.statemachine.config.builders.StateMachineStateConfigurer;
+                public class ComplexStateConfig extends StateMachineConfigurerAdapter {
+                    @Override
+                    public void configure(StateMachineStateConfigurer states) throws Exception {
+                        states
+                            .withStates()
+                            .initial("CREATED")
+                            .fork("PROCESSING")
+                            .and()
+                            .withStates()
+                            .parent("PROCESSING")
+                            .region("SHIPPING")
+                            .and()
+                            .withStates()
+                            .parent("SHIPPING")
+                            .initial("PREPARING")
+                            .end("SHIPPED")
+                            .and()
+                            .withStates()
+                            .parent("PROCESSING")
+                            .region("BILLING")
+                            .and()
+                            .withStates()
+                            .parent("BILLING")
+                            .initial("CALCULATING")
+                            .end("BILLED")
+                            .and()
+                            .withStates()
+                            .parent("PROCESSING")
+                            .end("COMPLETED");
+                    }
+                }
+                """;
+
+        writeClass("ComplexStateConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("ComplexStateConfig");
+        aggregator.aggregateStates(td);
+
+        assertThat(aggregator.getInitialStates()).containsExactly("CREATED");
+        assertThat(aggregator.getEndStates()).containsExactlyInAnyOrder("COMPLETED", "SHIPPED", "BILLED");
+    }
+
+    @Test
+    void shouldFollowMethodCallsInOtherClasses() throws IOException {
+        String otherClassSource = """
+                package com.example;
+                public class ExternalTransitions {
+                    public static void addCommon(StateMachineTransitionConfigurer transitions) {
+                        transitions.withExternal().source("S1").target("S2").event("EXT_E1");
+                    }
+                }
+                """;
+        String mainSource = """
+                package com.example;
+                public class MainConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        ExternalTransitions.addCommon(transitions);
+                    }
+                }
+                """;
+
+        writeClass("ExternalTransitions", otherClassSource);
+        writeClass("MainConfig", mainSource);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("MainConfig");
+        List<Transition> transitions = aggregator.aggregateTransitions(td);
+
+        // Currently StateMachineAggregator only searches within the hierarchy of the searchStartClass.
+        // It might not find ExternalTransitions.addCommon unless it's static or we handle imports.
+        // Let's see if it works or if we need to improve the aggregator.
+        assertThat(transitions).hasSize(1);
+        assertThat(transitions.get(0).getEvent()).isEqualTo("EXT_E1");
+    }
+
+    @Test
+    void shouldHandleRecursionSafety() throws IOException {
+        String source = """
+                package com.example;
+                public class RecursiveConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        transitions.withExternal().source("S1").target("S2").event("E1");
+                        recursiveMethod(transitions);
+                    }
+                    private void recursiveMethod(StateMachineTransitionConfigurer transitions) {
+                        transitions.withExternal().source("S2").target("S3").event("E2");
+                        recursiveMethod(transitions);
+                    }
+                }
+                """;
+
+        writeClass("RecursiveConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("RecursiveConfig");
+        List<Transition> transitions = aggregator.aggregateTransitions(td);
+
+        // Should not crash and should find both transitions
+        assertThat(transitions).hasSize(2);
+        assertThat(transitions).anyMatch(t -> t.getEvent().equals("E1"));
+        assertThat(transitions).anyMatch(t -> t.getEvent().equals("E2"));
+    }
+
+    @Test
+    void shouldParseTransitionsInsideTryBlock() throws IOException {
+        String source = """
+                package com.example;
+                public class TryConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        try {
+                            transitions.withExternal().source("S1").target("S2").event("TRY_E1");
+                        } catch (Exception e) {
+                            // ignore
+                        }
+                    }
+                }
+                """;
+
+        writeClass("TryConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("TryConfig");
+        List<Transition> transitions = aggregator.aggregateTransitions(td);
+
+        assertThat(transitions).hasSize(1);
+        assertThat(transitions.get(0).getEvent()).isEqualTo("TRY_E1");
+    }
+
     private void writeClass(String name, String source) throws IOException {
         Files.writeString(tempDir.resolve(name + ".java"), source);
     }

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

@@ -1,392 +0,0 @@
-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");
-    }
-} 

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

@@ -1,344 +0,0 @@
-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");
-    }
-} 

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/common/CodebaseContextTest.java

@@ -97,4 +97,53 @@ class CodebaseContextTest {
         assertThat(entryPoints).hasSize(1);
         assertThat(entryPoints.get(0).getName().getIdentifier()).isEqualTo("LeafConfig");
     }
+
+    @Test
+    void shouldFindBeanMethodsReturningStateMachine() throws IOException {
+        String source = """
+                package com.example;
+                import org.springframework.context.annotation.Bean;
+                import org.springframework.statemachine.StateMachine;
+                public class MyApp {
+                    @Bean
+                    public StateMachine<String, String> stateMachine() {
+                        return null;
+                    }
+                    
+                    public StateMachine<String, String> notABean() {
+                        return null;
+                    }
+                }
+                """;
+        Files.writeString(tempDir.resolve("MyApp.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        var beanMethods = context.findBeanMethodsReturning(java.util.Set.of("StateMachine"));
+        assertThat(beanMethods).hasSize(1);
+        assertThat(beanMethods.get(0).getName().getIdentifier()).isEqualTo("stateMachine");
+    }
+
+    @Test
+    void shouldFindBeanMethodsReturningStateMachineWithQualifiedName() throws IOException {
+        String source = """
+                package com.example;
+                import org.springframework.context.annotation.Bean;
+                public class MyApp {
+                    @Bean
+                    public org.springframework.statemachine.StateMachine stateMachine() {
+                        return null;
+                    }
+                }
+                """;
+        Files.writeString(tempDir.resolve("MyApp.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        var beanMethods = context.findBeanMethodsReturning(java.util.Set.of("StateMachine"));
+        assertThat(beanMethods).hasSize(1);
+        assertThat(beanMethods.get(0).getName().getIdentifier()).isEqualTo("stateMachine");
+    }
 }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/in/AstFileFinderTest.java

@@ -1,314 +0,0 @@
-package click.kamil.springstatemachineexporter.ast.in;
-
-import org.junit.jupiter.api.Test;
-
-import java.util.Collections;
-import java.util.List;
-import java.util.Set;
-
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
-
-class AstFileFinderTest {
-
-    private static final String CLASS_WITH_ANNOTATION = """
-            @Service
-            public class TestService {
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            """;
-
-    private static final String CLASS_WITH_MULTIPLE_ANNOTATIONS = """
-            @Service
-            @Component
-            public class TestService {
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            """;
-
-    private static final String INTERFACE_WITH_ANNOTATION = """
-            @FunctionalInterface
-            public interface TestInterface {
-                void doSomething();
-            }
-            """;
-
-    private static final String CLASS_WITHOUT_ANNOTATION = """
-            public class TestService {
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            """;
-
-    private static final String CLASS_WITH_DIFFERENT_ANNOTATION = """
-            @Repository
-            public class TestRepository {
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            """;
-
-    private static final String CLASS_WITH_FULLY_QUALIFIED_ANNOTATION = """
-            @org.springframework.stereotype.Service
-            public class TestService {
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            """;
-
-    private static final String MULTIPLE_CLASSES = """
-            @Service
-            public class TestService {
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            
-            public class AnotherClass {
-                // no annotation
-            }
-            """;
-
-    private static final String CLASS_WITH_ANNOTATION_AND_FIELDS = """
-            @Service
-            public class TestService {
-                private String name;
-            
-                public void doSomething() {
-                    // implementation
-                }
-            }
-            """;
-
-    private static final String EMPTY_SOURCE = "";
-
-    private static final String INVALID_JAVA_SYNTAX = """
-            @Service
-            public class TestService {
-                // missing closing brace
-            """;
-
-    @Test
-    void shouldReturnTrue_whenClassHasMatchingAnnotation() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION, annotationNames);
-
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnTrue_whenClassHasMultipleMatchingAnnotations() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service", "Component");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_MULTIPLE_ANNOTATIONS, annotationNames);
-
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnFalse_whenClassHasNoAnnotation() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITHOUT_ANNOTATION, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnFalse_whenClassHasDifferentAnnotation() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_DIFFERENT_ANNOTATION, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnTrue_whenClassHasFullyQualifiedAnnotation() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("org.springframework.stereotype.Service");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_FULLY_QUALIFIED_ANNOTATION, annotationNames);
-
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnTrue_whenMultipleClassesAndOneHasAnnotation() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(MULTIPLE_CLASSES, annotationNames);
-
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnFalse_whenInterfaceHasAnnotation() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("FunctionalInterface");
-
-        boolean result = finder.hasClassWithAnnotation(INTERFACE_WITH_ANNOTATION, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnFalse_whenEmptyAnnotationList() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = Collections.emptyList();
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnFalse_whenEmptySource() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(EMPTY_SOURCE, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldThrowException_whenNullAnnotationList() {
-        AstFileFinder finder = new AstFileFinder();
-
-        // when & then
-        assertThatThrownBy(() ->
-                finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION, null)
-        ).isInstanceOf(NullPointerException.class);
-    }
-
-    @Test
-    void shouldReturnTrue_whenClassHasAnnotationAndFields() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION_AND_FIELDS, annotationNames);
-
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldHandleInvalidJavaSyntax() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(INVALID_JAVA_SYNTAX, annotationNames);
-
-        // AST parser might still create a valid AST even with syntax errors
-        // The result depends on how the parser handles the invalid syntax
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnTrue_whenAnnotationNameMatchesCaseInsensitive() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("service"); // lowercase
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION, annotationNames);
-
-        assertThat(result).isFalse(); // AST parser is case-sensitive
-    }
-
-    @Test
-    void shouldReturnTrue_whenMultipleAnnotationNamesProvided() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Component", "Service", "Repository");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION, annotationNames);
-
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnFalse_whenNoMatchingAnnotationNames() {
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Component", "Repository", "Controller");
-
-        boolean result = finder.hasClassWithAnnotation(CLASS_WITH_ANNOTATION, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnFalse_whenSourceContainsOnlyComments() {
-        String sourceWithComments = """
-                // This is a comment
-                /* Another comment */
-                /** JavaDoc comment */
-                """;
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(sourceWithComments, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnFalse_whenSourceContainsOnlyPackageDeclaration() {
-        String sourceWithPackage = "package com.example;";
-        AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
-
-        boolean result = finder.hasClassWithAnnotation(sourceWithPackage, annotationNames);
-
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnTrue_whenFunctionReturnsTargetType() {
-        String source = """
-                public class Test {
-                    @Bean
-                    public StateMachine<S, E> build() { return null; }
-                }
-                """;
-        AstFileFinder finder = new AstFileFinder();
-        boolean result = finder.hasFunctionReturningType(source, Set.of("StateMachine"));
-        assertThat(result).isTrue();
-    }
-
-    @Test
-    void shouldReturnFalse_whenNoFunctionReturnsTargetType() {
-        String source = """
-                public class Test {
-                    @Bean
-                    public void build() {}
-                }
-                """;
-        AstFileFinder finder = new AstFileFinder();
-        boolean result = finder.hasFunctionReturningType(source, Set.of("StateMachine"));
-        assertThat(result).isFalse();
-    }
-
-    @Test
-    void shouldReturnTrue_whenFunctionReturnsQualifiedTargetType() {
-        String source = """
-                public class Test {
-                    @Bean
-                    public org.springframework.statemachine.StateMachine build() { return null; }
-                }
-                """;
-        AstFileFinder finder = new AstFileFinder();
-        boolean result = finder.hasFunctionReturningType(source, Set.of("StateMachine"));
-        assertThat(result).isTrue();
-    }
-}