Fix agent message type checks

4 files changed, 13 insertions, 4 deletions

diff --git a/src/main/java/pl/edu/mimuw/cloudatlas/agent/EventBus.java b/src/main/java/pl/edu/mimuw/cloudatlas/agent/EventBus.java
index b84d693..37c125a 100644
--- a/src/main/java/pl/edu/mimuw/cloudatlas/agent/EventBus.java
+++ b/src/main/java/pl/edu/mimuw/cloudatlas/agent/EventBus.java
@@ -43,13 +43,13 @@ public class EventBus implements Runnable {
     }
 
     public void routeMessage(AgentMessage msg) throws InterruptedException {
-        assert msg.correctType == msg.getDestinationModule();
+        assert msg.getCorrectMessageType() == msg.getDestinationModule();
         System.out.println("Event bus routing message");
         executors.get(msg.getDestinationModule()).addMessage(msg);
     }
 
     public void addMessage(AgentMessage msg) throws InterruptedException {
-        assert msg.correctType == msg.getDestinationModule();
+        assert msg.getCorrectMessageType() == msg.getDestinationModule();
         this.events.put(msg);
     }
 }
diff --git a/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/AgentMessage.java b/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/AgentMessage.java
index ba15a7a..b71831f 100644
--- a/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/AgentMessage.java
+++ b/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/AgentMessage.java
@@ -1,7 +1,6 @@
 package pl.edu.mimuw.cloudatlas.agent.messages;
 
-public class AgentMessage {
-    public AgentModule correctType;
+public abstract class AgentMessage {
 
     public enum AgentModule {
         TIMER_SCHEDULER,
@@ -30,6 +29,8 @@ public class AgentMessage {
         this.timestamp = System.currentTimeMillis() / 1000L;
     }
 
+    public abstract AgentModule getCorrectMessageType();
+
     public String getMessageId() {
         return messageId;
     }
diff --git a/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/TimerSchedulerMessage.java b/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/TimerSchedulerMessage.java
index 23ec700..9a558b7 100644
--- a/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/TimerSchedulerMessage.java
+++ b/src/main/java/pl/edu/mimuw/cloudatlas/agent/messages/TimerSchedulerMessage.java
@@ -43,4 +43,9 @@ public class TimerSchedulerMessage extends AgentMessage {
     public String getRequestId() { return requestId; }
 
     public void setRequestId(String requestId) { this.requestId = requestId; }
+
+    @Override
+    public AgentModule getCorrectMessageType() {
+        return AgentModule.TIMER_SCHEDULER;
+    }
 }
diff --git a/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/TimerScheduler.java b/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/TimerScheduler.java
index 921f985..a08d060 100644
--- a/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/TimerScheduler.java
+++ b/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/TimerScheduler.java
@@ -13,18 +13,21 @@ import java.util.Timer;
  * Tasks declared as inherited from TimerTask
  *
  * TODO: add request id and custom time
+ * TODO: enable messaging from scheduled tasks
  */
 public class TimerScheduler extends Module {
     private Timer timer;
 
     public TimerScheduler(AgentMessage.AgentModule moduleType) {
         super(moduleType);
+        assert moduleType == AgentMessage.AgentModule.TIMER_SCHEDULER;
         this.timer = new Timer();
         System.out.println("TimerScheduler instance initialized");
     }
 
     @Override
     public void handle(AgentMessage event) throws InterruptedException {
+        assert event.getDestinationModule() == event.getCorrectMessageType();
         TimerSchedulerMessage timerEvent = (TimerSchedulerMessage) event;
         this.timer.schedule(timerEvent.getTask(), timerEvent.getDelay());
         System.out.println("Task with delay " + timerEvent.getDelay() + " scheduled");