8
8
9
9
package org .opensearch .search .sandboxing .cancellation ;
10
10
11
- import org .junit .Before ;
12
- import org .mockito .Mock ;
13
- import org .mockito .Mockito ;
14
- import org .mockito .MockitoAnnotations ;
15
11
import org .opensearch .cluster .metadata .Sandbox ;
16
12
import org .opensearch .search .sandboxing .SandboxLevelResourceUsageView ;
17
13
import org .opensearch .search .sandboxing .resourcetype .SandboxResourceType ;
18
14
import org .opensearch .test .OpenSearchTestCase ;
15
+ import org .junit .Before ;
19
16
20
17
import java .util .Collections ;
21
18
import java .util .HashMap ;
24
21
import java .util .Map ;
25
22
import java .util .Set ;
26
23
24
+ import org .mockito .Mock ;
25
+ import org .mockito .Mockito ;
26
+ import org .mockito .MockitoAnnotations ;
27
+
27
28
import static org .mockito .Mockito .mock ;
28
29
import static org .mockito .Mockito .when ;
29
30
30
31
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