From efdff7b9f143d340bd27ac35a28658209fda47d5 Mon Sep 17 00:00:00 2001
From: Fangjin Yang <fangjin@metamarketsgroup.com>
Date: Fri, 25 Jan 2013 14:33:05 -0800
Subject: [PATCH] fix unit test bug

---
 .../druid/merger/coordinator/RemoteTaskRunner.java   | 10 ++++++++++
 .../merger/coordinator/RemoteTaskRunnerTest.java     | 12 +++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/merger/src/main/java/com/metamx/druid/merger/coordinator/RemoteTaskRunner.java b/merger/src/main/java/com/metamx/druid/merger/coordinator/RemoteTaskRunner.java
index 067356ceddc..fa6a7185058 100644
--- a/merger/src/main/java/com/metamx/druid/merger/coordinator/RemoteTaskRunner.java
+++ b/merger/src/main/java/com/metamx/druid/merger/coordinator/RemoteTaskRunner.java
@@ -264,6 +264,16 @@ public class RemoteTaskRunner implements TaskRunner
     return zkWorkers.size();
   }
 
+  public boolean isTaskRunning(String taskId)
+  {
+    for (WorkerWrapper workerWrapper : zkWorkers.values()) {
+      if (workerWrapper.getRunningTasks().contains(taskId)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   @Override
   public void run(Task task, TaskContext context, TaskCallback callback)
   {
diff --git a/merger/src/test/java/com/metamx/druid/merger/coordinator/RemoteTaskRunnerTest.java b/merger/src/test/java/com/metamx/druid/merger/coordinator/RemoteTaskRunnerTest.java
index 6f90029e6b1..4f4bc9fdfed 100644
--- a/merger/src/test/java/com/metamx/druid/merger/coordinator/RemoteTaskRunnerTest.java
+++ b/merger/src/test/java/com/metamx/druid/merger/coordinator/RemoteTaskRunnerTest.java
@@ -196,8 +196,13 @@ public class RemoteTaskRunnerTest
     );
 
     // Really don't like this way of waiting for the task to appear
-    while (remoteTaskRunner.getNumWorkers() == 0) {
+    int count = 0;
+    while (!remoteTaskRunner.isTaskRunning("task1")) {
       Thread.sleep(500);
+      if (count > 10) {
+        throw new ISE("WTF?! Task still not announced in ZK?");
+      }
+      count++;
     }
 
     final MutableBoolean callbackCalled = new MutableBoolean(false);
@@ -364,8 +369,13 @@ public class RemoteTaskRunnerTest
         String.format("%s/worker1", announcementsPath),
         jsonMapper.writeValueAsBytes(worker1)
     );
+    int count = 0;
     while (remoteTaskRunner.getNumWorkers() == 0) {
       Thread.sleep(500);
+      if (count > 10) {
+        throw new ISE("WTF?! Still can't find worker!");
+      }
+      count++;
     }
   }