From 66c3f3c5ea829c801bb97fcf0b46aa2d96478f28 Mon Sep 17 00:00:00 2001
From: maobaolong <baoloongmao@tencent.com>
Date: Mon, 6 Jan 2025 15:18:03 +0800
Subject: [PATCH] [MINOR] test: fix NPE for test
 PartitionBalanceAssignmentStrategyTest (#2322)

### What changes were proposed in this pull request?

Fix NPE for test PartitionBalanceAssignmentStrategyTest

### Why are the changes needed?

It will NPE if The metrics have not register to MetricsSystem yet for UTs.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Existing UT.
---
 .../assignment/BasicAssignmentStrategyTest.java     | 13 +++++++++++++
 .../PartitionBalanceAssignmentStrategyTest.java     | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/BasicAssignmentStrategyTest.java b/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/BasicAssignmentStrategyTest.java
index fd03f19d5a..36afe77eef 100644
--- a/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/BasicAssignmentStrategyTest.java
+++ b/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/BasicAssignmentStrategyTest.java
@@ -28,12 +28,15 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import org.apache.hadoop.conf.Configuration;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.apache.uniffle.common.PartitionRange;
 import org.apache.uniffle.coordinator.CoordinatorConf;
 import org.apache.uniffle.coordinator.ServerNode;
 import org.apache.uniffle.coordinator.SimpleClusterManager;
+import org.apache.uniffle.coordinator.metric.CoordinatorMetrics;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -45,6 +48,16 @@ public class BasicAssignmentStrategyTest {
   Set<String> tags = Sets.newHashSet("test");
   private int shuffleNodesMax = 7;
 
+  @BeforeEach
+  public void setUp() {
+    CoordinatorMetrics.register();
+  }
+
+  @AfterEach
+  public void clear() {
+    CoordinatorMetrics.clear();
+  }
+
   @Test
   public void testAssign() throws Exception {
     CoordinatorConf ssc = new CoordinatorConf();
diff --git a/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/PartitionBalanceAssignmentStrategyTest.java b/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/PartitionBalanceAssignmentStrategyTest.java
index 60d6cb873f..6c6c591e75 100644
--- a/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/PartitionBalanceAssignmentStrategyTest.java
+++ b/coordinator/src/test/java/org/apache/uniffle/coordinator/strategy/assignment/PartitionBalanceAssignmentStrategyTest.java
@@ -30,11 +30,14 @@
 import com.google.common.collect.Sets;
 import com.google.common.util.concurrent.Uninterruptibles;
 import org.apache.hadoop.conf.Configuration;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.apache.uniffle.coordinator.CoordinatorConf;
 import org.apache.uniffle.coordinator.ServerNode;
 import org.apache.uniffle.coordinator.SimpleClusterManager;
+import org.apache.uniffle.coordinator.metric.CoordinatorMetrics;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -46,6 +49,16 @@ public class PartitionBalanceAssignmentStrategyTest {
   private int shuffleNodesMax = 5;
   private Set<String> tags = Sets.newHashSet("test");
 
+  @BeforeEach
+  public void setUp() {
+    CoordinatorMetrics.register();
+  }
+
+  @AfterEach
+  public void clear() {
+    CoordinatorMetrics.clear();
+  }
+
   @Test
   public void testAssign() throws Exception {
     CoordinatorConf ssc = new CoordinatorConf();