From c30e1e49305ce1e4c074e473355dc263e555b6f9 Mon Sep 17 00:00:00 2001
From: Jason Gustafson <jason@confluent.io>
Date: Tue, 28 May 2019 16:49:39 -0700
Subject: [PATCH] MINOR: Fix transient failure in
 PartitionTest.testAddAndRemoveMetrics

---
 core/src/test/scala/unit/kafka/cluster/PartitionTest.scala | 3 +++
 .../test/scala/unit/kafka/network/SocketServerTest.scala   | 3 +--
 .../unit/kafka/server/AbstractFetcherManagerTest.scala     | 4 ++--
 .../unit/kafka/server/AbstractFetcherThreadTest.scala      | 3 +--
 core/src/test/scala/unit/kafka/utils/TestUtils.scala       | 7 +++++++
 5 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
index b3e5ade009a55..2b12a04cf4d4b 100644
--- a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
+++ b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
@@ -66,6 +66,8 @@ class PartitionTest {
 
   @Before
   def setup(): Unit = {
+    TestUtils.clearYammerMetrics()
+
     val logProps = createLogProperties(Map.empty)
     logConfig = LogConfig(logProps)
 
@@ -104,6 +106,7 @@ class PartitionTest {
   def tearDown(): Unit = {
     logManager.shutdown()
     Utils.delete(tmpDir)
+    TestUtils.clearYammerMetrics()
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
index 0a8ce1d0455eb..5162d7e53f7c2 100644
--- a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
+++ b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
@@ -68,8 +68,7 @@ class SocketServerTest {
   val localAddress = InetAddress.getLoopbackAddress
 
   // Clean-up any metrics left around by previous tests
-  for (metricName <- YammerMetrics.defaultRegistry.allMetrics.keySet.asScala)
-    YammerMetrics.defaultRegistry.removeMetric(metricName)
+  TestUtils.clearYammerMetrics()
 
   val server = new SocketServer(config, metrics, Time.SYSTEM, credentialProvider)
   server.startup()
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
index ec57a0fc847bd..15ce97133809d 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
@@ -19,6 +19,7 @@ package kafka.server
 import com.yammer.metrics.Metrics
 import com.yammer.metrics.core.Gauge
 import kafka.cluster.BrokerEndPoint
+import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
 import org.easymock.EasyMock
 import org.junit.{Before, Test}
@@ -30,8 +31,7 @@ class AbstractFetcherManagerTest {
 
   @Before
   def cleanMetricRegistry(): Unit = {
-    for (metricName <- Metrics.defaultRegistry().allMetrics().keySet().asScala)
-      Metrics.defaultRegistry().removeMetric(metricName)
+    TestUtils.clearYammerMetrics()
   }
 
   private def getMetricValue(name: String): Any = {
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
index 1d8249a715298..55c38a1deab80 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
@@ -52,8 +52,7 @@ class AbstractFetcherThreadTest {
 
   @Before
   def cleanMetricRegistry(): Unit = {
-    for (metricName <- Metrics.defaultRegistry().allMetrics().keySet().asScala)
-      Metrics.defaultRegistry().removeMetric(metricName)
+    TestUtils.clearYammerMetrics()
   }
 
   private def allMetricsNames: Set[String] = Metrics.defaultRegistry().allMetrics().asScala.keySet.map(_.getName)
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index c7f5c24f14c7d..ea8d2b3e549fa 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -35,6 +35,7 @@ import kafka.security.auth.{Acl, Authorizer, Resource}
 import kafka.server._
 import kafka.server.checkpoints.OffsetCheckpointFile
 import Implicits._
+import com.yammer.metrics.Metrics
 import kafka.controller.LeaderIsrAndControllerEpoch
 import kafka.zk._
 import org.apache.kafka.clients.CommonClientConfigs
@@ -1483,4 +1484,10 @@ object TestUtils extends Logging {
       .foldLeft(0.0)((total, metric) => total + metric.metricValue.asInstanceOf[Double])
     total.toLong
   }
+
+  def clearYammerMetrics(): Unit = {
+    for (metricName <- Metrics.defaultRegistry.allMetrics.keySet.asScala)
+      Metrics.defaultRegistry.removeMetric(metricName)
+  }
+
 }