From 7077214c81930a906804c6fdbf589f83817a85c9 Mon Sep 17 00:00:00 2001 From: Kamil Patryk Kozakowski Date: Fri, 19 Jun 2026 20:56:14 +0200 Subject: [PATCH] better enricher --- .../enricher/TransitionLinkerEnricher.java | 34 +++++++++++-- .../service/GenericEventDetector.java | 3 +- .../TransitionLinkerEnricherTest.java | 26 ++++++++++ .../service/CallGraphBuilderTest.java | 50 +++++++++++++++++++ .../service/GenericEventDetectorTest.java | 29 +++++++++++ 5 files changed, 136 insertions(+), 6 deletions(-) diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricher.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricher.java index 059e06e..a7e1eda 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricher.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricher.java @@ -45,6 +45,17 @@ public class TransitionLinkerEnricher implements AnalysisEnricher { boolean isWildcard = triggerEvent.equals("event") || triggerEvent.equals("e") || triggerEvent.equals("msg") || triggerEvent.equals("message") || triggerEvent.equals("payload") || triggerEvent.matches(".*\\.get[A-Z].*\\(\\)"); + + if (isWildcard) { + String targetVar = chain.getContextMachineId(); + if (targetVar == null && chain.getTriggerPoint() != null) { + targetVar = chain.getTriggerPoint().getStateMachineId(); + } + if (targetVar == null || targetVar.isEmpty()) { + // Wildcard without specific routing context is too broad, do not link + continue; + } + } if (isWildcard || smEvent.equals(triggerEvent) || triggerEvent.toLowerCase().contains(smEvent.toLowerCase()) || smEvent.toLowerCase().contains(triggerEvent.toLowerCase())) { // Event matches or is a wildcard @@ -52,17 +63,27 @@ public class TransitionLinkerEnricher implements AnalysisEnricher { String smSourceRaw = smSourceState.fullIdentifier() != null ? smSourceState.fullIdentifier() : smSourceState.rawName(); String smSource = simplify(smSourceRaw); if (triggerSource == null || triggerSource.equals(smSource)) { - for (State smTargetState : t.getTargetStates()) { - String sourceRaw = smSourceState.fullIdentifier() != null ? smSourceState.fullIdentifier() : smSourceState.rawName(); - String targetRaw = smTargetState.fullIdentifier() != null ? smTargetState.fullIdentifier() : smTargetState.rawName(); + if (t.getTargetStates() == null || t.getTargetStates().isEmpty()) { MatchedTransition mt = MatchedTransition.builder() - .sourceState(sourceRaw) - .targetState(targetRaw) + .sourceState(smSourceRaw) + .targetState(smSourceRaw) .event(smEventRaw) .build(); if (isRoutedToCorrectMachine(chain, result.getName())) { matched.add(mt); } + } else { + for (State smTargetState : t.getTargetStates()) { + String targetRaw = smTargetState.fullIdentifier() != null ? smTargetState.fullIdentifier() : smTargetState.rawName(); + MatchedTransition mt = MatchedTransition.builder() + .sourceState(smSourceRaw) + .targetState(targetRaw) + .event(smEventRaw) + .build(); + if (isRoutedToCorrectMachine(chain, result.getName())) { + matched.add(mt); + } + } } } } @@ -115,6 +136,9 @@ public class TransitionLinkerEnricher implements AnalysisEnricher { if (name == null) return null; // Strip common suffixes String simplified = name.replaceAll("(?i)(.*)(Event|Action|Transition|Command)(s)?$", "$1"); + if (simplified.isEmpty()) { + simplified = name; + } // Simplify full identifiers to just the last part (enum name) simplified = simplified.replaceAll("^.*\\.([A-Z0-9_]+)$", "$1"); // For remaining full caps with underscores (like EVENT_X), keep as is or try to simplify diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetector.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetector.java index 282fc7e..baec1fe 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetector.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetector.java @@ -302,7 +302,8 @@ public class GenericEventDetector { } else if (current.arguments().isEmpty() && current.getExpression() instanceof SimpleName sn) { // If the event is obtained by calling a method on a provider/supplier parameter, // return the provider's name so CallGraphBuilder can trace the lambda argument - return sn.getIdentifier(); + String traced = extractEventFromMessageBuilder(sn); + return traced != null ? traced : sn.getIdentifier(); } Expression receiver = current.getExpression(); diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricherTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricherTest.java index e4f87fd..ba91d4f 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricherTest.java +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricherTest.java @@ -128,6 +128,7 @@ class TransitionLinkerEnricherTest { CallChain chain = CallChain.builder() .triggerPoint(TriggerPoint.builder().event("event").build()) + .contextMachineId("testMachine") .build(); AnalysisResult result = AnalysisResult.builder() @@ -151,6 +152,7 @@ class TransitionLinkerEnricherTest { CallChain chain = CallChain.builder() .triggerPoint(TriggerPoint.builder().event("richEvent.getType()").build()) + .contextMachineId("testMachine") .build(); AnalysisResult result = AnalysisResult.builder() @@ -164,4 +166,28 @@ class TransitionLinkerEnricherTest { CallChain updatedChain = result.getMetadata().getCallChains().get(0); assertThat(updatedChain.getMatchedTransitions()).hasSize(1); } + + @Test + void shouldIgnoreWildcardWhenContextMachineIdIsMissing() { + Transition t1 = new Transition(); + t1.setSourceStates(List.of(State.of("NEW", "NEW"))); + t1.setTargetStates(List.of(State.of("PAID", "PAID"))); + t1.setEvent(Event.of("PAY", "PAY")); + + CallChain chain = CallChain.builder() + .triggerPoint(TriggerPoint.builder().event("event").build()) + // Intentionally NO contextMachineId or stateMachineId + .build(); + + AnalysisResult result = AnalysisResult.builder() + .name("testMachine") + .transitions(List.of(t1)) + .metadata(CodebaseMetadata.builder().callChains(List.of(chain)).build()) + .build(); + + enricher.enrich(result, null, null); + + CallChain updatedChain = result.getMetadata().getCallChains().get(0); + assertThat(updatedChain.getMatchedTransitions()).isNullOrEmpty(); + } } diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CallGraphBuilderTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CallGraphBuilderTest.java index 30722a2..7d4fe72 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CallGraphBuilderTest.java +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CallGraphBuilderTest.java @@ -115,4 +115,54 @@ class CallGraphBuilderTest { // Since it's tracing constant concatenation: "machine:" + orderId, which might resolve dynamically or string format. // It should at least be present (it typically extracts 'machine:' + orderId into expression). } + + @Test + void shouldResolveTriggerPointParametersAcrossChain() throws IOException { + String source = """ + package com.example; + public class OrderController { + private OrderService service; + public void createOrder() { + service.updateOrderState(OrderEvents.CREATE); + } + } + + class OrderService { + private StateMachine sm; + public void updateOrderState(OrderEvents event) { + sm.sendEvent(event); + } + } + + class StateMachine { + public void sendEvent(Object event) {} + } + + enum OrderEvents { CREATE, PAY } + """; + Path tempDir = Files.createTempDirectory("callgraph_test"); + Files.writeString(tempDir.resolve("OrderConfig.java"), source); + + CodebaseContext context = new CodebaseContext(); + context.scan(tempDir); + + CallGraphBuilder builder = new CallGraphBuilder(context); + + EntryPoint entryPoint = EntryPoint.builder() + .className("com.example.OrderController") + .methodName("createOrder") + .build(); + + TriggerPoint trigger = TriggerPoint.builder() + .className("com.example.OrderService") + .methodName("updateOrderState") + .event("event") + .build(); + + List chains = builder.findChains(List.of(entryPoint), List.of(trigger)); + + assertThat(chains).hasSize(1); + CallChain chain = chains.get(0); + assertThat(chain.getTriggerPoint().getEvent()).isEqualTo("OrderEvents.CREATE"); + } } diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetectorTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetectorTest.java index b79ba7b..1175ee1 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetectorTest.java +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/GenericEventDetectorTest.java @@ -45,6 +45,35 @@ public class GenericEventDetectorTest { assertThat(triggers.get(0).getMethodName()).isEqualTo("a"); } + @Test + void testMessageBuilderIntermediateVariable(@TempDir Path tempDir) throws IOException { + Path dir = tempDir.resolve("com/example"); + Files.createDirectories(dir); + Files.writeString(dir.resolve("MyService.java"), + "package com.example;\n" + + "import org.springframework.messaging.support.MessageBuilder;\n" + + "import org.springframework.messaging.Message;\n" + + "import org.springframework.statemachine.StateMachine;\n" + + "public class MyService {\n" + + " private StateMachine stateMachine;\n" + + " private void a(MyInterface myEvent) {\n" + + " MessageBuilder messageBuilder = MessageBuilder.withPayload(myEvent).setHeader(\"k\", \"v\");\n" + + " stateMachine.sendEvent(messageBuilder.build());\n" + + " }\n" + + "}\n"); + + CodebaseContext context = new CodebaseContext(); + context.scan(tempDir); + + GenericEventDetector detector = new GenericEventDetector(context, new ConstantResolver(), List.of()); + CompilationUnit cu = context.getCompilationUnits().iterator().next(); + List triggers = detector.detect(cu); + + assertThat(triggers).hasSize(1); + assertThat(triggers.get(0).getEvent()).isEqualTo("myEvent"); + assertThat(triggers.get(0).getMethodName()).isEqualTo("a"); + } + @Test void testEnumGetterUnionExtraction(@TempDir Path tempDir) throws IOException { Path dir = tempDir.resolve("com/example");