diff --git a/src/main/java/teetime/framework/AbstractStage.java b/src/main/java/teetime/framework/AbstractStage.java
index ad20973adae4bd1065369ffe1a57c71acd52ef46..21ea950557186bda73c0fa6c3be438e13d3a5065 100644
--- a/src/main/java/teetime/framework/AbstractStage.java
+++ b/src/main/java/teetime/framework/AbstractStage.java
@@ -60,8 +60,9 @@ public abstract class AbstractStage extends Stage {
 		Class<? extends ISignal> signalClass = signal.getClass();
 
 		Set<InputPort<?>> signalReceivedInputPorts;
-		signalReceivedInputPorts = signalMap.get(signalClass);
-		if (signalReceivedInputPorts == null) {
+		if (signalMap.containsKey(signalClass)) {
+			signalReceivedInputPorts = signalMap.get(signalClass);
+		} else {
 			signalReceivedInputPorts = new HashSet<InputPort<?>>();
 			signalMap.put(signalClass, signalReceivedInputPorts);
 		}