diff --git a/core/src/test/scala/unit/kafka/admin/AdminTest.scala b/core/src/test/scala/unit/kafka/admin/AdminTest.scala
index e1f68dd..7a31b51 100644
--- a/core/src/test/scala/unit/kafka/admin/AdminTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AdminTest.scala
@@ -374,35 +374,38 @@ class AdminTest extends JUnit3Suite with ZooKeeperTestHarness {
     var controllerId = ZkUtils.getController(zkClient)
     var controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
     var partitionsRemaining = controller.shutdownBroker(2)
-    assertEquals(0, partitionsRemaining)
-    var topicMetadata = AdminUtils.fetchTopicMetadataFromZk(topic, zkClient)
-    var leaderAfterShutdown = topicMetadata.partitionsMetadata.head.leader.get.id
-    assertTrue(leaderAfterShutdown != leaderBeforeShutdown)
-    assertEquals(2, topicMetadata.partitionsMetadata.head.isr.size)
-
-    leaderBeforeShutdown = leaderAfterShutdown
-    controllerId = ZkUtils.getController(zkClient)
-    controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
-    partitionsRemaining = controller.shutdownBroker(1)
-    assertEquals(0, partitionsRemaining)
-    topicMetadata = AdminUtils.fetchTopicMetadataFromZk(topic, zkClient)
-    leaderAfterShutdown = topicMetadata.partitionsMetadata.head.leader.get.id
-    assertTrue(leaderAfterShutdown != leaderBeforeShutdown)
-    assertEquals(1, topicMetadata.partitionsMetadata.head.isr.size)
-
-    leaderBeforeShutdown = leaderAfterShutdown
-    controllerId = ZkUtils.getController(zkClient)
-    controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
-    partitionsRemaining = controller.shutdownBroker(0)
-    assertEquals(1, partitionsRemaining)
-    topicMetadata = AdminUtils.fetchTopicMetadataFromZk(topic, zkClient)
-    leaderAfterShutdown = topicMetadata.partitionsMetadata.head.leader.get.id
-    assertTrue(leaderAfterShutdown == leaderBeforeShutdown)
-    assertEquals(1, topicMetadata.partitionsMetadata.head.isr.size)
-
-    servers.foreach(_.shutdown())
-
-
+    try {
+      assertEquals(0, partitionsRemaining)
+      var topicMetadata = AdminUtils.fetchTopicMetadataFromZk(topic, zkClient)
+      var leaderAfterShutdown = topicMetadata.partitionsMetadata.head.leader.get.id
+      assertTrue(leaderAfterShutdown != leaderBeforeShutdown)
+      // assertEquals(2, topicMetadata.partitionsMetadata.head.isr.size)
+      assertEquals(2, controller.controllerContext.allLeaders(TopicAndPartition("test", 1)).leaderAndIsr.isr.size)
+
+      leaderBeforeShutdown = leaderAfterShutdown
+      controllerId = ZkUtils.getController(zkClient)
+      controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
+      partitionsRemaining = controller.shutdownBroker(1)
+      assertEquals(0, partitionsRemaining)
+      topicMetadata = AdminUtils.fetchTopicMetadataFromZk(topic, zkClient)
+      leaderAfterShutdown = topicMetadata.partitionsMetadata.head.leader.get.id
+      assertTrue(leaderAfterShutdown != leaderBeforeShutdown)
+      // assertEquals(1, topicMetadata.partitionsMetadata.head.isr.size)
+      assertEquals(1, controller.controllerContext.allLeaders(TopicAndPartition("test", 1)).leaderAndIsr.isr.size)
+
+      leaderBeforeShutdown = leaderAfterShutdown
+      controllerId = ZkUtils.getController(zkClient)
+      controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
+      partitionsRemaining = controller.shutdownBroker(0)
+      assertEquals(1, partitionsRemaining)
+      topicMetadata = AdminUtils.fetchTopicMetadataFromZk(topic, zkClient)
+      leaderAfterShutdown = topicMetadata.partitionsMetadata.head.leader.get.id
+      assertTrue(leaderAfterShutdown == leaderBeforeShutdown)
+      assertEquals(1, controller.controllerContext.allLeaders(TopicAndPartition("test", 1)).leaderAndIsr.isr.size)
+    }
+    finally {
+      servers.foreach(_.shutdown())
+    }
   }
 
   private def checkIfReassignPartitionPathExists(): Boolean = {