Skip to content

Commit 07e4a51

Browse files
kiranprakash154kiranprakash154
committed
Some tests and stubs
Signed-off-by: Kiran Prakash <[email protected]>
1 parent 5e17561 commit 07e4a51

File tree

3 files changed

+229
-3
lines changed

3 files changed

+229
-3
lines changed

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,16 +52,24 @@ public class ResourceLimit {
5252
/*
5353
Stub class only
5454
*/
55+
SandboxResourceType resourceType;
56+
Long threshold;
57+
58+
ResourceLimit(SandboxResourceType resourceType, Long threshold) {
59+
this.resourceType = resourceType;
60+
this.threshold = threshold;
61+
}
62+
5563
public Long getThresholdInLong() {
56-
return 0L;
64+
return threshold;
5765
}
5866

5967
public SandboxResourceType getResourceType() {
60-
return null;
68+
return resourceType;
6169
}
6270

6371
public Long getThreshold() {
64-
return 0L;
72+
return threshold;
6573
}
6674
}
6775

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.sandboxing;
10+
11+
import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType;
12+
import org.opensearch.tasks.Task;
13+
import org.opensearch.test.OpenSearchTestCase;
14+
15+
import java.util.List;
16+
import java.util.Map;
17+
18+
import static org.opensearch.search.sandboxing.cancellation.SandboxCancellationStrategyTestHelpers.getRandomTask;
19+
20+
public class SandboxLevelResourceUsageViewTest extends OpenSearchTestCase {
21+
Map<SandboxResourceType, Long> resourceUsage;
22+
List<Task> activeTasks;
23+
24+
public void setUp() throws Exception {
25+
super.setUp();
26+
resourceUsage = Map.of(
27+
SandboxResourceType.fromString("JVM"), 34L,
28+
SandboxResourceType.fromString("CPU"), 12L
29+
);
30+
activeTasks = List.of(getRandomTask(4321));
31+
}
32+
33+
public void testGetResourceUsageData() {
34+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView(
35+
"1234",
36+
resourceUsage,
37+
activeTasks
38+
);
39+
Map<SandboxResourceType, Long> resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData();
40+
assertTrue(assertResourceUsageData(resourceUsageData));
41+
}
42+
43+
public void testGetResourceUsageDataDefault() {
44+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234");
45+
Map<SandboxResourceType, Long> resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData();
46+
assertTrue(resourceUsageData.isEmpty());
47+
}
48+
49+
public void testGetActiveTasks() {
50+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView(
51+
"1234",
52+
resourceUsage,
53+
activeTasks
54+
);
55+
List<Task> activeTasks = sandboxLevelResourceUsageView.getActiveTasks();
56+
assertEquals(1, activeTasks.size());
57+
assertEquals(4321, activeTasks.get(0).getId());
58+
}
59+
60+
public void testGetActiveTasksDefault() {
61+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234");
62+
List<Task> activeTasks = sandboxLevelResourceUsageView.getActiveTasks();
63+
assertTrue(activeTasks.isEmpty());
64+
}
65+
66+
private boolean assertResourceUsageData(Map<SandboxResourceType, Long> resourceUsageData) {
67+
return resourceUsageData.get(SandboxResourceType.fromString("JVM")) == 34L &&
68+
resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L;
69+
}
70+
}

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

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.sandboxing.cancellation;
10+
11+
import org.junit.Before;
12+
import org.mockito.Mock;
13+
import org.mockito.Mockito;
14+
import org.mockito.MockitoAnnotations;
15+
import org.opensearch.cluster.metadata.Sandbox;
16+
import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView;
17+
import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType;
18+
import org.opensearch.test.OpenSearchTestCase;
19+
20+
import java.util.Collections;
21+
import java.util.HashMap;
22+
import java.util.HashSet;
23+
import java.util.List;
24+
import java.util.Map;
25+
import java.util.Set;
26+
27+
import static org.mockito.Mockito.mock;
28+
import static org.mockito.Mockito.when;
29+
30+
public class DefaultTaskCancellationTest extends OpenSearchTestCase {
31+
@Mock
32+
private TaskSelectionStrategy mockStrategy;
33+
34+
@Mock
35+
private SandboxLevelResourceUsageView mockView;
36+
37+
private Map<String, SandboxLevelResourceUsageView> sandboxLevelViews;
38+
private Set<Sandbox> activeSandboxes;
39+
40+
@Before
41+
public void setup() {
42+
MockitoAnnotations.openMocks(this);
43+
sandboxLevelViews = new HashMap<>();
44+
activeSandboxes = new HashSet<>();
45+
}
46+
47+
public void testConstructor() {
48+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
49+
assertNotNull(cancellation);
50+
}
51+
52+
public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() {
53+
// setup mocks for sandbox1
54+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
55+
// setup mocks for sandbox2
56+
Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L);
57+
// add both sandboxes to active sandboxes
58+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
59+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
60+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
61+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
62+
63+
// only sandbox1 should be returned as it is breaching the threshold
64+
assertEquals(1, result.size());
65+
assertTrue(result.contains(sandbox1));
66+
assertFalse(result.contains(sandbox2));
67+
}
68+
69+
public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() {
70+
// setup mocks for sandbox1
71+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
72+
// setup mocks for sandbox2
73+
Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L);
74+
// add both sandboxes to active sandboxes
75+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
76+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
77+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
78+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
79+
80+
// only sandbox1 should be returned as it is breaching the threshold
81+
assertEquals(1, result.size());
82+
assertTrue(result.contains(sandbox1));
83+
assertFalse(result.contains(sandbox2));
84+
}
85+
86+
public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() {
87+
// setup mocks for sandbox1
88+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
89+
// setup mocks for sandbox2
90+
Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L);
91+
// add both sandboxes to active sandboxes
92+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
93+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
94+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
95+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
96+
97+
// Both sandboxes should be returned as it is breaching the threshold
98+
assertEquals(2, result.size());
99+
assertTrue(result.contains(sandbox1));
100+
assertTrue(result.contains(sandbox2));
101+
}
102+
103+
public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() {
104+
// setup mocks for sandbox1
105+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
106+
// setup mocks for sandbox2
107+
Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L);
108+
// add both sandboxes to active sandboxes
109+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
110+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
111+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
112+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
113+
114+
// Both sandboxes should be returned as it is breaching the threshold
115+
assertEquals(2, result.size());
116+
assertTrue(result.contains(sandbox1));
117+
assertTrue(result.contains(sandbox2));
118+
}
119+
120+
// Utility methods
121+
private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) {
122+
Sandbox sandbox = Mockito.mock(Sandbox.class);
123+
when(sandbox.getId()).thenReturn(id);
124+
125+
Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold);
126+
when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock));
127+
128+
SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage);
129+
sandboxLevelViews.put(id, mockView);
130+
131+
return sandbox;
132+
}
133+
134+
private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) {
135+
Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class);
136+
SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr);
137+
when(resourceLimitMock.getResourceType()).thenReturn(resourceType);
138+
when(resourceLimitMock.getThreshold()).thenReturn(threshold);
139+
return resourceLimitMock;
140+
}
141+
142+
private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) {
143+
SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class);
144+
SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr);
145+
when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage));
146+
return mockView;
147+
}
148+
}

0 commit comments

Comments
 (0)