Fix extended state and @OnTransition handling

- Add better support for enums with @OnTransition fixes #25
- Add enum handling with extended state fixes #26

spring-statemachine-core/src/main/java/org/springframework/statemachine/ExtendedState.java

@@ -21,17 +21,17 @@ import java.util.Map;
  * Extended states are used to supplement state machine with a variables. If
  * extended state is used a complete condition of a state machine is a
  * combination of its state an extended state variables.
- * 
+ *
  * @author Janne Valkealahti
  *
  */
 public interface ExtendedState {
-	
+
 	/**
 	 * Gets the extended state variables.
 	 *
 	 * @return the extended state variables
 	 */
-	Map<String, Object> getVariables();
+	Map<Object, Object> getVariables();
 
 }

spring-statemachine-core/src/main/java/org/springframework/statemachine/StateContext.java

@@ -37,6 +37,15 @@ public interface StateContext<S, E> {
 	 */
 	MessageHeaders getMessageHeaders();
 
+	/**
+	 * Gets the message header. If header is not a {@code String} object's
+	 * {@link Object#toString()} method is used to resolve a key name.
+	 *
+	 * @param header the header
+	 * @return the message header
+	 */
+	Object getMessageHeader(Object header);
+
 	/**
 	 * Gets the state machine extended state.
 	 *

spring-statemachine-core/src/main/java/org/springframework/statemachine/support/AbstractStateMachine.java

@@ -56,6 +56,7 @@ import org.springframework.statemachine.trigger.TimerTrigger;
 import org.springframework.statemachine.trigger.Trigger;
 import org.springframework.statemachine.trigger.TriggerListener;
 import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
 
 /**
  * Base implementation of a {@link StateMachine} loosely modelled from UML state
@@ -277,9 +278,6 @@ public abstract class AbstractStateMachine<S, E> extends LifecycleObjectSupport
 	private void switchToState(State<S,E> state, Message<E> event, Transition<S,E> transition) {
 		exitFromState(currentState, event, transition);
 		notifyStateChanged(currentState, state);
-
-		callHandlers(currentState, state, event);
-
 		setCurrentState(state, event, transition);
 
 		// TODO: should handle triggerles transition some how differently
@@ -411,6 +409,7 @@ public abstract class AbstractStateMachine<S, E> extends LifecycleObjectSupport
 			}
 
 			notifyTransitionStart(transition);
+			callHandlers(transition.getSource(), transition.getTarget(), queuedEvent);
 			boolean transit = transition.transit(stateContext);
 			if (transit && transition.getKind() != TransitionKind.INTERNAL) {
 				switchToState(transition.getTarget(), queuedEvent, transition);
@@ -424,16 +423,12 @@ public abstract class AbstractStateMachine<S, E> extends LifecycleObjectSupport
 	}
 
 	private void callHandlers(State<S,E> sourceState, State<S,E> targetState, Message<E> event) {
-
-
-
 		if (sourceState != null && targetState != null) {
 			MessageHeaders messageHeaders = event != null ? event.getHeaders() : new MessageHeaders(
 					new HashMap<String, Object>());
 			StateContext<S, E> stateContext = new DefaultStateContext<S, E>(messageHeaders, extendedState, null);
 			getStateMachineHandlerResults(getStateMachineHandlers(sourceState, targetState), stateContext);
 		}
-
 	}
 
 	private List<Object> getStateMachineHandlerResults(List<StateMachineHandler<S, E>> stateMachineHandlers, final StateContext<S, E> stateContext) {
@@ -475,9 +470,28 @@ public abstract class AbstractStateMachine<S, E> extends LifecycleObjectSupport
 			String source = annotation.source();
 			String target = annotation.target();
 			// TODO: need major fixes
-			String s = sourceState.getIds().iterator().next().toString();
-			String t = targetState.getIds().iterator().next().toString();
-			if (s.equals(source) && t.equals(target)) {
+			boolean handle = false;
+			if (StringUtils.hasText(source) && StringUtils.hasText(target)) {
+				if (StateMachineUtils.containsAtleastOneEqualString(
+						StateMachineUtils.toStringCollection(sourceState.getIds()), source)
+						&& StateMachineUtils.containsAtleastOneEqualString(
+								StateMachineUtils.toStringCollection(targetState.getIds()), target)) {
+					handle = true;
+				}
+			} else if (StringUtils.hasText(source)) {
+				if (StateMachineUtils.containsAtleastOneEqualString(
+						StateMachineUtils.toStringCollection(sourceState.getIds()), source)) {
+					handle = true;
+				}
+			} else if (StringUtils.hasText(target)) {
+				if (StateMachineUtils.containsAtleastOneEqualString(
+						StateMachineUtils.toStringCollection(targetState.getIds()), target)) {
+					handle = true;
+				}
+
+			}
+			System.out.println("XXX " + source + "/" + target + "/" +handle);
+			if (handle) {
 				handlersList.add(entry.getValue());
 			}
 		}

spring-statemachine-core/src/main/java/org/springframework/statemachine/support/DefaultExtendedState.java

@@ -22,23 +22,23 @@ import org.springframework.statemachine.ExtendedState;
 
 /**
  * Default implementation of a {@link ExtendedState}.
- * 
+ *
  * @author Janne Valkealahti
  *
  */
 public class DefaultExtendedState implements ExtendedState {
 
-	private final Map<String, Object> variables;
-	
+	private final Map<Object, Object> variables;
+
 	/**
 	 * Instantiates a new default extended state.
 	 */
 	public DefaultExtendedState() {
-		this.variables = new HashMap<String, Object>();
+		this.variables = new HashMap<Object, Object>();
 	}
-	
+
 	@Override
-	public Map<String, Object> getVariables() {
+	public Map<Object, Object> getVariables() {
 		return variables;
 	}
 

spring-statemachine-core/src/main/java/org/springframework/statemachine/support/DefaultStateContext.java

@@ -21,11 +21,11 @@ import org.springframework.statemachine.StateContext;
 import org.springframework.statemachine.transition.Transition;
 
 public class DefaultStateContext<S, E> implements StateContext<S, E> {
-	
+
 	private final MessageHeaders messageHeaders;
-	
+
 	private final ExtendedState extendedState;
-	
+
 	private final Transition<S,E> transition;
 
 	public DefaultStateContext(MessageHeaders messageHeaders, ExtendedState extendedState, Transition<S,E> transition) {
@@ -39,11 +39,21 @@ public class DefaultStateContext<S, E> implements StateContext<S, E> {
 		return messageHeaders;
 	}
 
+	@Override
+	public Object getMessageHeader(Object header) {
+		if (header instanceof String) {
+			return messageHeaders.get((String)header);
+		} else if (header instanceof Enum<?>) {
+			return messageHeaders.get(((Enum<?>)header).toString());
+		}
+		return null;
+	}
+
 	@Override
 	public ExtendedState getExtendedState() {
 		return extendedState;
 	}
-	
+
 	@Override
 	public Transition<S, E> getTransition() {
 		return transition;

spring-statemachine-core/src/main/java/org/springframework/statemachine/support/StateMachineUtils.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.statemachine.support;
 
+import java.util.ArrayList;
 import java.util.Collection;
 
 /**
@@ -46,4 +47,30 @@ public abstract class StateMachineUtils {
 		return false;
 	}
 
+	public static <S> Collection<String> toStringCollection(Collection<S> collection) {
+		Collection<String> c = new ArrayList<String>();
+		for (S item : collection) {
+			c.add(item.toString());
+		}
+		return c;
+	}
+
+	public static boolean containsAtleastOneEqualString(Collection<String> left, String right) {
+		Collection<String> r = new ArrayList<String>(1);
+		r.add(right);
+		return containsAtleastOneEqualString(left, r);
+	}
+
+	public static boolean containsAtleastOneEqualString(Collection<String> left, Collection<String> right) {
+		if (left == null || right == null) {
+			return false;
+		}
+		for (String id : left) {
+			if (right.contains(id)) {
+				return true;
+			}
+		}
+		return false;
+	}
+
 }