Correctly mock new 3 param get method

Signed-off-by: owenhalpert <[email protected]>
opensearch-project · Jan 24, 2025 · 55280c5 · 55280c5
1 parent 5eb55ed
commit 55280c5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java b/src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java
@@ -157,6 +157,7 @@ public static void setUpClass() throws Exception {
         final NativeMemoryCacheManager nativeMemoryCacheManager = mock(NativeMemoryCacheManager.class);
         final NativeMemoryAllocation nativeMemoryAllocation = mock(NativeMemoryAllocation.class);
         when(nativeMemoryCacheManager.get(any(), anyBoolean())).thenReturn(nativeMemoryAllocation);
+        when(nativeMemoryCacheManager.get(any(), anyBoolean(), anyBoolean())).thenReturn(nativeMemoryAllocation);
 
         nativeMemoryCacheManagerMockedStatic.when(NativeMemoryCacheManager::getInstance).thenReturn(nativeMemoryCacheManager);