From 1d1e923446b3323cca89b9c783985850b77d8f53 Mon Sep 17 00:00:00 2001
From: Kaushal Kumar <ravi.kaushal97@gmail.com>
Date: Tue, 30 Jul 2024 12:43:38 -0700
Subject: [PATCH] fix UT to remove the verify for final method

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
---
 .../wlm/WorkloadManagementTransportRequestHandlerTests.java    | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportRequestHandlerTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportRequestHandlerTests.java
index 23a0040747b12..47a88bb9f24a5 100644
--- a/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportRequestHandlerTests.java
+++ b/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportRequestHandlerTests.java
@@ -49,8 +49,7 @@ public void testMessageReceivedForSearchWorkload() throws Exception {
         QueryGroupTask spyTask = getSpyTask();
 
         sut.messageReceived(request, mock(TransportChannel.class), spyTask);
-
-        verify(spyTask, times(1)).setQueryGroupId(threadPool.getThreadContext());
+        assertTrue(sut.isSearchWorkloadRequest(spyTask));
     }
 
     public void testMessageReceivedForNonSearchWorkload() throws Exception {