/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.heartbeat;

import com.google.common.collect.ImmutableList;
import org.junit.Assert;
import org.junit.Test;

/**
 * Unit tests for {@link HeartbeatContext}.
 */
public class HeartbeatContextTest {
  @Test
  public void allThreadsUseSleepingTimer() {
    for (String threadName : HeartbeatContext.getTimerClasses().keySet()) {
      Class<? extends HeartbeatTimer> timerClass = HeartbeatContext.getTimerClass(threadName);
      Assert.assertTrue(timerClass.isAssignableFrom(SleepingTimer.class));
    }
  }

  @Test
  public void canUpdateToScheduledTimer() throws Exception {
    try (ManuallyScheduleHeartbeat.Resource h = new ManuallyScheduleHeartbeat.Resource(
        ImmutableList.of(HeartbeatContext.MASTER_CHECKPOINT_SCHEDULING))) {
      Class<? extends HeartbeatTimer> timerClass =
          HeartbeatContext.getTimerClass(HeartbeatContext.MASTER_CHECKPOINT_SCHEDULING);
      Assert.assertTrue(timerClass.isAssignableFrom(ScheduledTimer.class));
    }
  }
}

@@ -1059,14 +1059,25 @@
  can
-Update
+TemporarilySwitch
 ToSc
@@ -1157,16 +1157,24 @@
 urce h =
+%0A       
  new Man
@@ -1205,25 +1205,16 @@
 esource(
-%0A        
 Immutabl
@@ -1243,184 +1243,172 @@
 ext.
-MASTER_CHECKPOINT_SCHEDULING))) %7B%0A      Class%3C? extends HeartbeatTimer%3E timerClass =%0A          HeartbeatContext.getTimerClass(HeartbeatContext.MASTER_CHECKPOINT_SCHEDULING
+WORKER_CLIENT))) %7B%0A      Assert.assertTrue(HeartbeatContext.getTimerClass(HeartbeatContext.WORKER_CLIENT)%0A          .isAssignableFrom(ScheduledTimer.class)
 );%0A  
+  %7D%0A
     
@@ -1421,34 +1421,95 @@
 .assertTrue(
-timerClass
+HeartbeatContext.getTimerClass(HeartbeatContext.WORKER_CLIENT)%0A        
 .isAssignabl
@@ -1511,32 +1511,31 @@
 gnableFrom(S
-cheduled
+leeping
 Timer.class)
@@ -1533,24 +1533,18 @@
 er.class));%0A
-    %7D%0A
   %7D%0A%7D%0A

Improve test