Some tests and stubs

Signed-off-by: Kiran Prakash <awskiran@amazon.com>

server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java

@@ -52,16 +52,24 @@ public class ResourceLimit {
         /*
         Stub class only
         */
+        SandboxResourceType resourceType;
+        Long threshold;
+
+        ResourceLimit(SandboxResourceType resourceType, Long threshold) {
+            this.resourceType = resourceType;
+            this.threshold = threshold;
+        }
+
         public Long getThresholdInLong() {
-            return 0L;
+            return threshold;
         }
 
         public SandboxResourceType getResourceType() {
-            return null;
+            return resourceType;
         }
 
         public Long getThreshold() {
-            return 0L;
+            return threshold;
         }
     }
 

server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java

@@ -0,0 +1,70 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.sandboxing;
+
+import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType;
+import org.opensearch.tasks.Task;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.opensearch.search.sandboxing.cancellation.SandboxCancellationStrategyTestHelpers.getRandomTask;
+
+public class SandboxLevelResourceUsageViewTest extends OpenSearchTestCase {
+  Map<SandboxResourceType, Long> resourceUsage;
+  List<Task> activeTasks;
+
+  public void setUp() throws Exception {
+    super.setUp();
+    resourceUsage = Map.of(
+        SandboxResourceType.fromString("JVM"), 34L,
+        SandboxResourceType.fromString("CPU"), 12L
+    );
+    activeTasks = List.of(getRandomTask(4321));
+  }
+
+  public void testGetResourceUsageData() {
+    SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView(
+        "1234",
+        resourceUsage,
+        activeTasks
+    );
+    Map<SandboxResourceType, Long> resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData();
+    assertTrue(assertResourceUsageData(resourceUsageData));
+  }
+
+  public void testGetResourceUsageDataDefault() {
+    SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234");
+    Map<SandboxResourceType, Long> resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData();
+    assertTrue(resourceUsageData.isEmpty());
+  }
+
+  public void testGetActiveTasks() {
+    SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView(
+        "1234",
+        resourceUsage,
+        activeTasks
+    );
+    List<Task> activeTasks = sandboxLevelResourceUsageView.getActiveTasks();
+    assertEquals(1, activeTasks.size());
+    assertEquals(4321, activeTasks.get(0).getId());
+  }
+
+  public void testGetActiveTasksDefault() {
+    SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234");
+    List<Task> activeTasks = sandboxLevelResourceUsageView.getActiveTasks();
+    assertTrue(activeTasks.isEmpty());
+  }
+
+  private boolean assertResourceUsageData(Map<SandboxResourceType, Long> resourceUsageData) {
+    return resourceUsageData.get(SandboxResourceType.fromString("JVM")) == 34L &&
+        resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L;
+  }
+}

server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java

@@ -0,0 +1,148 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.sandboxing.cancellation;
+
+import org.junit.Before;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.opensearch.cluster.metadata.Sandbox;
+import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView;
+import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class DefaultTaskCancellationTest extends OpenSearchTestCase {
+  @Mock
+  private TaskSelectionStrategy mockStrategy;
+
+  @Mock
+  private SandboxLevelResourceUsageView mockView;
+
+  private Map<String, SandboxLevelResourceUsageView> sandboxLevelViews;
+  private Set<Sandbox> activeSandboxes;
+
+  @Before
+  public void setup() {
+    MockitoAnnotations.openMocks(this);
+    sandboxLevelViews = new HashMap<>();
+    activeSandboxes = new HashSet<>();
+  }
+
+  public void testConstructor() {
+    DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
+    assertNotNull(cancellation);
+  }
+
+  public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() {
+    // setup mocks for sandbox1
+    Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
+    // setup mocks for sandbox2
+    Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L);
+    // add both sandboxes to active sandboxes
+    Collections.addAll(activeSandboxes, sandbox1, sandbox2);
+    // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
+    DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
+    List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
+
+    // only sandbox1 should be returned as it is breaching the threshold
+    assertEquals(1, result.size());
+    assertTrue(result.contains(sandbox1));
+    assertFalse(result.contains(sandbox2));
+  }
+
+  public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() {
+    // setup mocks for sandbox1
+    Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
+    // setup mocks for sandbox2
+    Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L);
+    // add both sandboxes to active sandboxes
+    Collections.addAll(activeSandboxes, sandbox1, sandbox2);
+    // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
+    DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
+    List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
+
+    // only sandbox1 should be returned as it is breaching the threshold
+    assertEquals(1, result.size());
+    assertTrue(result.contains(sandbox1));
+    assertFalse(result.contains(sandbox2));
+  }
+
+  public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() {
+    // setup mocks for sandbox1
+    Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
+    // setup mocks for sandbox2
+    Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L);
+    // add both sandboxes to active sandboxes
+    Collections.addAll(activeSandboxes, sandbox1, sandbox2);
+    // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
+    DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
+    List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
+
+    // Both sandboxes should be returned as it is breaching the threshold
+    assertEquals(2, result.size());
+    assertTrue(result.contains(sandbox1));
+    assertTrue(result.contains(sandbox2));
+  }
+
+  public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() {
+    // setup mocks for sandbox1
+    Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
+    // setup mocks for sandbox2
+    Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L);
+    // add both sandboxes to active sandboxes
+    Collections.addAll(activeSandboxes, sandbox1, sandbox2);
+    // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
+    DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
+    List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
+
+    // Both sandboxes should be returned as it is breaching the threshold
+    assertEquals(2, result.size());
+    assertTrue(result.contains(sandbox1));
+    assertTrue(result.contains(sandbox2));
+  }
+
+  // Utility methods
+  private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) {
+    Sandbox sandbox = Mockito.mock(Sandbox.class);
+    when(sandbox.getId()).thenReturn(id);
+
+    Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold);
+    when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock));
+
+    SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage);
+    sandboxLevelViews.put(id, mockView);
+
+    return sandbox;
+  }
+
+  private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) {
+    Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class);
+    SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr);
+    when(resourceLimitMock.getResourceType()).thenReturn(resourceType);
+    when(resourceLimitMock.getThreshold()).thenReturn(threshold);
+    return resourceLimitMock;
+  }
+
+  private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) {
+    SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class);
+    SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr);
+    when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage));
+    return mockView;
+  }
+}