fix more unit test failure

Author: wirybeaver

pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/replication/RealtimeSegmentCopierTest.java

@@ -22,6 +22,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import org.apache.hc.core5.http.ClassicHttpRequest;
 import org.apache.pinot.common.utils.SimpleHttpResponse;
 import org.apache.pinot.common.utils.http.HttpClient;
 import org.apache.pinot.controller.ControllerConf;
@@ -34,8 +35,6 @@
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.anyMap;
-import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -79,15 +78,16 @@ public void testCopy() throws Exception {
 
     doReturn(_pinotFS).when(_copier).getPinotFS(any(URI.class));
     when(_pinotFS.exists(any(URI.class))).thenReturn(false);
+    when(_pinotFS.copy(any(URI.class), any(URI.class))).thenReturn(true);
 
     SimpleHttpResponse response = mock(SimpleHttpResponse.class);
     when(response.getStatusCode()).thenReturn(200);
-    when(_httpClient.sendJsonPostRequest(any(URI.class), anyString(), anyMap())).thenReturn(response);
+    when(_httpClient.sendRequest(any(ClassicHttpRequest.class), any(Integer.class))).thenReturn(response);
 
     _copier.copy(tableNameWithType, segmentName, payload, metadata);
 
     verify(_pinotFS).copy(any(URI.class), any(URI.class));
-    verify(_httpClient).sendJsonPostRequest(any(URI.class), anyString(), anyMap());
+    verify(_httpClient).sendRequest(any(ClassicHttpRequest.class), any(Integer.class));
   }
 
   @Test
@@ -105,11 +105,11 @@ public void testCopyExisting() throws Exception {
 
     SimpleHttpResponse response = mock(SimpleHttpResponse.class);
     when(response.getStatusCode()).thenReturn(200);
-    when(_httpClient.sendJsonPostRequest(any(URI.class), anyString(), anyMap())).thenReturn(response);
+    when(_httpClient.sendRequest(any(ClassicHttpRequest.class), any(Integer.class))).thenReturn(response);
 
     _copier.copy(tableNameWithType, segmentName, payload, metadata);
 
     verify(_pinotFS, never()).copy(any(URI.class), any(URI.class));
-    verify(_httpClient).sendJsonPostRequest(any(URI.class), anyString(), anyMap());
+    verify(_httpClient).sendRequest(any(ClassicHttpRequest.class), any(Integer.class));
   }
 }

pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/replication/ZkBasedTableReplicationObserverTest.java

@@ -20,6 +20,7 @@
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
@@ -35,11 +36,13 @@
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyLong;
 import static org.mockito.Mockito.anyMap;
 import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class ZkBasedTableReplicationObserverTest {
 
@@ -59,12 +62,18 @@ public void tearDown() throws Exception {
   }
 
   @Test
-  public void testObserver() {
+  public void testObserver() throws Exception {
     String jobId = "job1";
     String tableName = "table1";
     List<String> segments = Arrays.asList("seg1", "seg2", "seg3");
     WatermarkInductionResult res = new WatermarkInductionResult(Collections.emptyList(), segments);
 
+    Map<String, String> baseMetadata = new HashMap<>();
+    baseMetadata.put(CommonConstants.ControllerJob.JOB_ID, jobId);
+    baseMetadata.put(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE, tableName);
+    when(_pinotHelixResourceManager.commonTableReplicationJobMetadata(eq(tableName), eq(jobId), anyLong(),
+        eq(res))).thenReturn(baseMetadata);
+
     ZkBasedTableReplicationObserver observer =
         new ZkBasedTableReplicationObserver(jobId, tableName, res, _pinotHelixResourceManager);
 
@@ -92,12 +101,18 @@ public void testObserver() {
   }
 
   @Test
-  public void testObserverError() {
+  public void testObserverError() throws Exception {
     String jobId = "job1";
     String tableName = "table1";
     List<String> segments = Arrays.asList("seg1");
     WatermarkInductionResult res = new WatermarkInductionResult(Collections.emptyList(), segments);
 
+    Map<String, String> baseMetadata = new HashMap<>();
+    baseMetadata.put(CommonConstants.ControllerJob.JOB_ID, jobId);
+    baseMetadata.put(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE, tableName);
+    when(_pinotHelixResourceManager.commonTableReplicationJobMetadata(eq(tableName), eq(jobId), anyLong(),
+        eq(res))).thenReturn(baseMetadata);
+
     ZkBasedTableReplicationObserver observer =
         new ZkBasedTableReplicationObserver(jobId, tableName, res, _pinotHelixResourceManager);