Clean up test

n1v0lg · Nov 5, 2023 · 146291f · 146291f
1 parent e736d19
commit 146291f
Showing 1 changed file with 70 additions and 50 deletions.
diff --git a/...va/org/elasticsearch/xpack/security/authc/support/mapper/NativeRoleMappingStoreTests.java b/...va/org/elasticsearch/xpack/security/authc/support/mapper/NativeRoleMappingStoreTests.java
@@ -200,56 +200,18 @@ public void testResolveRolesWithLastLoadCache() throws Exception {
         final ThreadPool mockThreadPool = mock(ThreadPool.class);
         when(mockThreadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY));
         when(client.threadPool()).thenReturn(mockThreadPool);
-        final SearchRequestBuilder searchRequestBuilder = Mockito.spy(new SearchRequestBuilder(client, SearchAction.INSTANCE));
-        when(client.prepareSearch(eq(SECURITY_MAIN_ALIAS))).thenReturn(searchRequestBuilder);
-        doAnswer(invocation -> {
-            @SuppressWarnings("unchecked")
-            final var listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
-            final var searchHit = new SearchHit(
-                randomIntBetween(0, Integer.MAX_VALUE),
-                NativeRoleMappingStore.getIdForName(mapping.getName())
-            );
-            try (XContentBuilder builder = JsonXContent.contentBuilder()) {
-                mapping.toXContent(builder, ToXContent.EMPTY_PARAMS);
-                searchHit.sourceRef(BytesReference.bytes(builder));
-            }
-            final var internalSearchResponse = new InternalSearchResponse(
-                new SearchHits(
-                    new SearchHit[] { searchHit },
-                    new TotalHits(1, TotalHits.Relation.EQUAL_TO),
-                    randomFloat(),
-                    null,
-                    null,
-                    null
-                ),
-                null,
-                null,
-                null,
-                false,
-                null,
-                0
-            );
-            final var searchResponse = new SearchResponse(
-                internalSearchResponse,
-                randomAlphaOfLengthBetween(3, 8),
-                1,
-                1,
-                0,
-                10,
-                null,
-                null
-            );
-            listener.onResponse(searchResponse);
-            return null;
-        }).when(client).search(any(SearchRequest.class), anyActionListener());
+        when(client.prepareSearch(eq(SECURITY_MAIN_ALIAS))).thenReturn(
+            Mockito.spy(new SearchRequestBuilder(client, SearchAction.INSTANCE))
+        );
+        doAnswerWithSearchResult(client, mapping);
 
-        final SecurityIndexManager securityIndex = mock(SecurityIndexManager.class);
         final ScriptService scriptService = new ScriptService(
             Settings.EMPTY,
             Collections.singletonMap(MustacheScriptEngine.NAME, new MustacheScriptEngine()),
             ScriptModule.CORE_CONTEXTS,
             () -> 1L
         );
+        final SecurityIndexManager securityIndex = mock(SecurityIndexManager.class);
         when(securityIndex.isAvailable(SecurityIndexManager.Availability.PRIMARY_SHARDS)).thenReturn(true);
         when(securityIndex.isAvailable(SecurityIndexManager.Availability.SEARCH_SHARDS)).thenReturn(true);
         when(securityIndex.indexExists()).thenReturn(true);
@@ -281,27 +243,85 @@ public void testResolveRolesWithLastLoadCache() throws Exception {
         // when security index is available, we still run a search
         verify(client, times(2)).search(any(SearchRequest.class), anyActionListener());
 
-        final boolean indexUnavailable = randomBoolean();
-        when(securityIndex.isAvailable(SecurityIndexManager.Availability.SEARCH_SHARDS)).thenReturn(indexUnavailable);
-        final boolean indexClosed = indexUnavailable || randomBoolean();
+        final boolean indexAvailable = randomBoolean();
+        when(securityIndex.isAvailable(SecurityIndexManager.Availability.SEARCH_SHARDS)).thenReturn(indexAvailable);
+        final boolean indexClosed = indexAvailable || randomBoolean();
         when(securityIndex.indexIsClosed()).thenReturn(indexClosed);
-        assertThat(resolveRoles(store, user), Matchers.containsInAnyOrder("role"));
+        assertThat(resolveRoles(store, user), Matchers.containsInAnyOrder(mapping.getRoles().toArray()));
         assertThat(store.getLastLoad(), contains(mapping));
         // index was unavailable so we returned result from cache; no new search
         verify(client, times(2)).search(any(SearchRequest.class), anyActionListener());
 
-        // When index does not exist we don't fetch from cache
+        // when index does not exist we don't fetch from cache
         when(securityIndex.indexExists()).thenReturn(false);
         assertThat(resolveRoles(store, user), is(empty()));
         assertThat(store.getLastLoad(), contains(mapping));
+
+        // new search result from index overwrites previous
+        when(securityIndex.indexExists()).thenReturn(true);
+        when(securityIndex.isAvailable(SecurityIndexManager.Availability.SEARCH_SHARDS)).thenReturn(true);
+        when(securityIndex.indexIsClosed()).thenReturn(false);
+        final ExpressionRoleMapping mapping2 = new ExpressionRoleMapping(
+            "mapping2",
+            new FieldExpression("dn", Collections.singletonList(new FieldValue("*"))),
+            List.of("role2"),
+            Collections.emptyList(),
+            Collections.emptyMap(),
+            true
+        );
+        doAnswerWithSearchResult(client, mapping2);
+        assertThat(resolveRoles(store, user), Matchers.containsInAnyOrder(mapping2.getRoles().toArray()));
+        assertThat(store.getLastLoad(), contains(mapping2));
+    }
+
+    private void doAnswerWithSearchResult(Client client, ExpressionRoleMapping mapping) {
+        doAnswer(invocation -> {
+            @SuppressWarnings("unchecked")
+            final var listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
+            final var searchHit = new SearchHit(
+                randomIntBetween(0, Integer.MAX_VALUE),
+                NativeRoleMappingStore.getIdForName(mapping.getName())
+            );
+            try (XContentBuilder builder = JsonXContent.contentBuilder()) {
+                mapping.toXContent(builder, ToXContent.EMPTY_PARAMS);
+                searchHit.sourceRef(BytesReference.bytes(builder));
+            }
+            final var internalSearchResponse = new InternalSearchResponse(
+                new SearchHits(
+                    new SearchHit[] { searchHit },
+                    new TotalHits(1, TotalHits.Relation.EQUAL_TO),
+                    randomFloat(),
+                    null,
+                    null,
+                    null
+                ),
+                null,
+                null,
+                null,
+                false,
+                null,
+                0
+            );
+            final var searchResponse = new SearchResponse(
+                internalSearchResponse,
+                randomAlphaOfLengthBetween(3, 8),
+                1,
+                1,
+                0,
+                10,
+                null,
+                null
+            );
+            listener.onResponse(searchResponse);
+            return null;
+        }).when(client).search(any(SearchRequest.class), anyActionListener());
     }
 
     private Set<String> resolveRoles(NativeRoleMappingStore store, UserRoleMapper.UserData user) throws InterruptedException,
         ExecutionException {
         final PlainActionFuture<Set<String>> future2 = new PlainActionFuture<>();
         store.resolveRoles(user, future2);
-        Set<String> actual = future2.get();
-        return actual;
+        return future2.get();
     }
 
     private String randomiseDn(String dn) {