Skip to content

Commit 98989ba

Browse files
spotless
Signed-off-by: Kiran Prakash <[email protected]>
1 parent 07e4a51 commit 98989ba

File tree

2 files changed

+162
-172
lines changed

2 files changed

+162
-172
lines changed

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

Lines changed: 39 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -18,53 +18,42 @@
1818
import static org.opensearch.search.sandboxing.cancellation.SandboxCancellationStrategyTestHelpers.getRandomTask;
1919

2020
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-
}
21+
Map<SandboxResourceType, Long> resourceUsage;
22+
List<Task> activeTasks;
23+
24+
public void setUp() throws Exception {
25+
super.setUp();
26+
resourceUsage = Map.of(SandboxResourceType.fromString("JVM"), 34L, SandboxResourceType.fromString("CPU"), 12L);
27+
activeTasks = List.of(getRandomTask(4321));
28+
}
29+
30+
public void testGetResourceUsageData() {
31+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234", resourceUsage, activeTasks);
32+
Map<SandboxResourceType, Long> resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData();
33+
assertTrue(assertResourceUsageData(resourceUsageData));
34+
}
35+
36+
public void testGetResourceUsageDataDefault() {
37+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234");
38+
Map<SandboxResourceType, Long> resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData();
39+
assertTrue(resourceUsageData.isEmpty());
40+
}
41+
42+
public void testGetActiveTasks() {
43+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234", resourceUsage, activeTasks);
44+
List<Task> activeTasks = sandboxLevelResourceUsageView.getActiveTasks();
45+
assertEquals(1, activeTasks.size());
46+
assertEquals(4321, activeTasks.get(0).getId());
47+
}
48+
49+
public void testGetActiveTasksDefault() {
50+
SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234");
51+
List<Task> activeTasks = sandboxLevelResourceUsageView.getActiveTasks();
52+
assertTrue(activeTasks.isEmpty());
53+
}
54+
55+
private boolean assertResourceUsageData(Map<SandboxResourceType, Long> resourceUsageData) {
56+
return resourceUsageData.get(SandboxResourceType.fromString("JVM")) == 34L
57+
&& resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L;
58+
}
59+
}

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

Lines changed: 123 additions & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,11 @@
88

99
package org.opensearch.search.sandboxing.cancellation;
1010

11-
import org.junit.Before;
12-
import org.mockito.Mock;
13-
import org.mockito.Mockito;
14-
import org.mockito.MockitoAnnotations;
1511
import org.opensearch.cluster.metadata.Sandbox;
1612
import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView;
1713
import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType;
1814
import org.opensearch.test.OpenSearchTestCase;
15+
import org.junit.Before;
1916

2017
import java.util.Collections;
2118
import java.util.HashMap;
@@ -24,125 +21,129 @@
2421
import java.util.Map;
2522
import java.util.Set;
2623

24+
import org.mockito.Mock;
25+
import org.mockito.Mockito;
26+
import org.mockito.MockitoAnnotations;
27+
2728
import static org.mockito.Mockito.mock;
2829
import static org.mockito.Mockito.when;
2930

3031
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-
}
32+
@Mock
33+
private TaskSelectionStrategy mockStrategy;
34+
35+
@Mock
36+
private SandboxLevelResourceUsageView mockView;
37+
38+
private Map<String, SandboxLevelResourceUsageView> sandboxLevelViews;
39+
private Set<Sandbox> activeSandboxes;
40+
41+
@Before
42+
public void setup() {
43+
MockitoAnnotations.openMocks(this);
44+
sandboxLevelViews = new HashMap<>();
45+
activeSandboxes = new HashSet<>();
46+
}
47+
48+
public void testConstructor() {
49+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
50+
assertNotNull(cancellation);
51+
}
52+
53+
public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() {
54+
// setup mocks for sandbox1
55+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
56+
// setup mocks for sandbox2
57+
Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L);
58+
// add both sandboxes to active sandboxes
59+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
60+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
61+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
62+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
63+
64+
// only sandbox1 should be returned as it is breaching the threshold
65+
assertEquals(1, result.size());
66+
assertTrue(result.contains(sandbox1));
67+
assertFalse(result.contains(sandbox2));
68+
}
69+
70+
public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() {
71+
// setup mocks for sandbox1
72+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
73+
// setup mocks for sandbox2
74+
Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L);
75+
// add both sandboxes to active sandboxes
76+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
77+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
78+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
79+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
80+
81+
// only sandbox1 should be returned as it is breaching the threshold
82+
assertEquals(1, result.size());
83+
assertTrue(result.contains(sandbox1));
84+
assertFalse(result.contains(sandbox2));
85+
}
86+
87+
public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() {
88+
// setup mocks for sandbox1
89+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
90+
// setup mocks for sandbox2
91+
Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L);
92+
// add both sandboxes to active sandboxes
93+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
94+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
95+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
96+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
97+
98+
// Both sandboxes should be returned as it is breaching the threshold
99+
assertEquals(2, result.size());
100+
assertTrue(result.contains(sandbox1));
101+
assertTrue(result.contains(sandbox2));
102+
}
103+
104+
public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() {
105+
// setup mocks for sandbox1
106+
Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L);
107+
// setup mocks for sandbox2
108+
Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L);
109+
// add both sandboxes to active sandboxes
110+
Collections.addAll(activeSandboxes, sandbox1, sandbox2);
111+
// create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom
112+
DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes);
113+
List<Sandbox> result = cancellation.getSandboxesToCancelFrom();
114+
115+
// Both sandboxes should be returned as it is breaching the threshold
116+
assertEquals(2, result.size());
117+
assertTrue(result.contains(sandbox1));
118+
assertTrue(result.contains(sandbox2));
119+
}
120+
121+
// Utility methods
122+
private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) {
123+
Sandbox sandbox = Mockito.mock(Sandbox.class);
124+
when(sandbox.getId()).thenReturn(id);
125+
126+
Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold);
127+
when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock));
128+
129+
SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage);
130+
sandboxLevelViews.put(id, mockView);
131+
132+
return sandbox;
133+
}
134+
135+
private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) {
136+
Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class);
137+
SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr);
138+
when(resourceLimitMock.getResourceType()).thenReturn(resourceType);
139+
when(resourceLimitMock.getThreshold()).thenReturn(threshold);
140+
return resourceLimitMock;
141+
}
142+
143+
private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) {
144+
SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class);
145+
SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr);
146+
when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage));
147+
return mockView;
148+
}
149+
}

0 commit comments

Comments
 (0)