1010import org .elasticsearch .action .ActionRequest ;
1111import org .elasticsearch .action .ActionResponse ;
1212import org .elasticsearch .action .ActionType ;
13+ import org .elasticsearch .action .admin .indices .refresh .RefreshRequest ;
14+ import org .elasticsearch .action .support .broadcast .BroadcastResponse ;
1315import org .elasticsearch .client .internal .Client ;
1416import org .elasticsearch .common .settings .Settings ;
1517import org .elasticsearch .common .util .concurrent .EsExecutors ;
2022import org .elasticsearch .test .ESTestCase ;
2123import org .elasticsearch .test .client .NoOpClient ;
2224import org .elasticsearch .threadpool .ThreadPool ;
25+ import org .elasticsearch .xpack .core .security .action .UpdateIndexMigrationVersionAction ;
2326import org .elasticsearch .xpack .core .security .action .UpdateIndexMigrationVersionResponse ;
2427import org .elasticsearch .xpack .core .security .support .SecurityMigrationTaskParams ;
2528import org .junit .Before ;
2629
30+ import java .util .List ;
2731import java .util .Map ;
2832import java .util .Set ;
2933import java .util .TreeMap ;
@@ -40,8 +44,11 @@ public class SecurityMigrationExecutorTests extends ESTestCase {
4044 private SecurityIndexManager securityIndexManager ;
4145
4246 private int updateIndexMigrationVersionActionInvocations ;
47+ private int refreshActionInvocations ;
4348
44- private boolean clientShouldThrowException = false ;
49+ private boolean updateVersionShouldThrowException = false ;
50+
51+ private boolean refreshIndexShouldThrowException = false ;
4552
4653 private AllocatedPersistentTask mockTask = mock (AllocatedPersistentTask .class );
4754
@@ -51,6 +58,7 @@ public void setUpMocks() {
5158 when (threadPool .getThreadContext ()).thenReturn (new ThreadContext (Settings .EMPTY ));
5259 when (threadPool .generic ()).thenReturn (EsExecutors .DIRECT_EXECUTOR_SERVICE );
5360 updateIndexMigrationVersionActionInvocations = 0 ;
61+ refreshActionInvocations = 0 ;
5462 client = new NoOpClient (threadPool ) {
5563 @ Override
5664 @ SuppressWarnings ("unchecked" )
@@ -59,12 +67,27 @@ protected <Request extends ActionRequest, Response extends ActionResponse> void
5967 Request request ,
6068 ActionListener <Response > listener
6169 ) {
62- if (clientShouldThrowException ) {
63- listener .onFailure (new IllegalStateException ("Bad client" ));
64- return ;
70+ if (request instanceof RefreshRequest ) {
71+ if (refreshIndexShouldThrowException ) {
72+ if (randomBoolean ()) {
73+ listener .onFailure (new IllegalStateException ("Refresh index failed" ));
74+ } else {
75+ listener .onResponse ((Response ) new BroadcastResponse (1 , 0 , 1 , List .of ()));
76+ }
77+ } else {
78+ refreshActionInvocations ++;
79+ listener .onResponse ((Response ) new BroadcastResponse (1 , 1 , 0 , List .of ()));
80+ }
81+ } else if (request instanceof UpdateIndexMigrationVersionAction .Request ) {
82+ if (updateVersionShouldThrowException ) {
83+ listener .onFailure (new IllegalStateException ("Update version failed" ));
84+ } else {
85+ updateIndexMigrationVersionActionInvocations ++;
86+ listener .onResponse ((Response ) new UpdateIndexMigrationVersionResponse ());
87+ }
88+ } else {
89+ fail ("Unexpected client request" );
6590 }
66- updateIndexMigrationVersionActionInvocations ++;
67- listener .onResponse ((Response ) new UpdateIndexMigrationVersionResponse ());
6891
6992 }
7093 };
@@ -85,6 +108,7 @@ public void testSuccessfulMigration() {
85108 verify (mockTask , times (1 )).markAsCompleted ();
86109 verify (mockTask , times (0 )).markAsFailed (any ());
87110 assertEquals (2 , updateIndexMigrationVersionActionInvocations );
111+ assertEquals (3 , refreshActionInvocations );
88112 assertEquals (2 , migrateInvocations [0 ]);
89113 }
90114
@@ -111,6 +135,7 @@ public void testNoMigrationMeetsRequirements() {
111135 verify (mockTask , times (1 )).markAsCompleted ();
112136 verify (mockTask , times (0 )).markAsFailed (any ());
113137 assertEquals (0 , updateIndexMigrationVersionActionInvocations );
138+ assertEquals (1 , refreshActionInvocations );
114139 assertEquals (0 , migrateInvocationsCounter [0 ]);
115140 }
116141
@@ -140,6 +165,7 @@ public void testPartialMigration() {
140165 securityMigrationExecutor .nodeOperation (mockTask , new SecurityMigrationTaskParams (0 , true ), mock (PersistentTaskState .class ));
141166 verify (mockTask , times (1 )).markAsCompleted ();
142167 verify (mockTask , times (0 )).markAsFailed (any ());
168+ assertEquals (3 , refreshActionInvocations );
143169 assertEquals (2 , updateIndexMigrationVersionActionInvocations );
144170 assertEquals (2 , migrateInvocations [0 ]);
145171 }
@@ -158,6 +184,7 @@ public void testNoMigrationNeeded() {
158184 verify (mockTask , times (1 )).markAsCompleted ();
159185 verify (mockTask , times (0 )).markAsFailed (any ());
160186 assertEquals (0 , updateIndexMigrationVersionActionInvocations );
187+ assertEquals (1 , refreshActionInvocations );
161188 assertEquals (0 , migrateInvocations [0 ]);
162189 }
163190
@@ -186,14 +213,13 @@ public int minMappingVersion() {
186213 }))
187214 );
188215
189- assertThrows (
190- IllegalStateException .class ,
191- () -> securityMigrationExecutor .nodeOperation (
216+ securityMigrationExecutor .nodeOperation (
192217 mockTask ,
193218 new SecurityMigrationTaskParams (0 , true ),
194219 mock (PersistentTaskState .class )
195- )
196- );
220+ );
221+ verify (mockTask , times (1 )).markAsFailed (any ());
222+ verify (mockTask , times (0 )).markAsCompleted ();
197223 }
198224
199225 public void testUpdateMigrationVersionThrowsException () {
@@ -205,12 +231,27 @@ public void testUpdateMigrationVersionThrowsException() {
205231 client ,
206232 new TreeMap <>(Map .of (1 , generateMigration (migrateInvocations , true ), 2 , generateMigration (migrateInvocations , true )))
207233 );
208- clientShouldThrowException = true ;
234+ updateVersionShouldThrowException = true ;
209235 securityMigrationExecutor .nodeOperation (mockTask , new SecurityMigrationTaskParams (0 , true ), mock (PersistentTaskState .class ));
210236 verify (mockTask , times (1 )).markAsFailed (any ());
211237 verify (mockTask , times (0 )).markAsCompleted ();
212238 }
213239
240+ public void testRefreshSecurityIndexThrowsException () {
241+ final int [] migrateInvocations = new int [1 ];
242+ SecurityMigrationExecutor securityMigrationExecutor = new SecurityMigrationExecutor (
243+ "test-task" ,
244+ threadPool .generic (),
245+ securityIndexManager ,
246+ client ,
247+ new TreeMap <>(Map .of (1 , generateMigration (migrateInvocations , true ), 2 , generateMigration (migrateInvocations , true )))
248+ );
249+ refreshIndexShouldThrowException = true ;
250+ securityMigrationExecutor .nodeOperation (mockTask , new SecurityMigrationTaskParams (0 , true ), mock (PersistentTaskState .class ));
251+ verify (mockTask , times (0 )).markAsFailed (any ());
252+ verify (mockTask , times (1 )).markAsCompleted ();
253+ }
254+
214255 private SecurityMigrations .SecurityMigration generateMigration (int [] migrateInvocationsCounter , boolean isEligible ) {
215256 SecurityMigrations .SecurityMigration migration = new SecurityMigrations .SecurityMigration () {
216257 @ Override
0 commit comments