Fix Failing Tests

Signed-off-by: Kiran Prakash <awskiran@amazon.com>
opensearch-project · May 30, 2024 · eb22efb · eb22efb
1 parent 225c31a
commit eb22efb
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java
@@ -15,4 +15,13 @@ public class CpuTimeResourceType extends SandboxResourceType {
   public long getResourceUsage(Task task) {
     return task.getTotalResourceStats().getCpuTimeInNanos();
   }
+  @Override
+  public boolean equals(Object obj) {
+    return obj instanceof CpuTimeResourceType;
+  }
+
+  @Override
+  public int hashCode() {
+    return "CPU".hashCode();
+  }
 }
diff --git a/...er/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java b/...er/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java
@@ -15,4 +15,13 @@ public class JvmMemoryResourceType extends SandboxResourceType {
   public long getResourceUsage(Task task) {
     return task.getTotalResourceStats().getMemoryInBytes();
   }
+  @Override
+  public boolean equals(Object obj) {
+    return obj instanceof JvmMemoryResourceType;
+  }
+
+  @Override
+  public int hashCode() {
+    return "JVM".hashCode();
+  }
 }
diff --git a/...org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java b/...org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java
@@ -13,20 +13,12 @@
 import org.opensearch.core.tasks.resourcetracker.ResourceStats;
 import org.opensearch.core.tasks.resourcetracker.ResourceStatsType;
 import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric;
-import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage;
-import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
-import org.opensearch.test.OpenSearchTestCase;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
 
-import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 import static org.opensearch.test.OpenSearchTestCase.randomLong;
 import static org.opensearch.test.OpenSearchTestCase.randomLongBetween;