Allow supplying a custom map to InMemoryOAuth2AuthorizedClientService

Author: vpavic

oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/InMemoryOAuth2AuthorizedClientService.java

@@ -36,8 +36,8 @@
  * @see Authentication
  */
 public final class InMemoryOAuth2AuthorizedClientService implements OAuth2AuthorizedClientService {
-	private final Map<OAuth2AuthorizedClientId, OAuth2AuthorizedClient> authorizedClients = new ConcurrentHashMap<>();
 	private final ClientRegistrationRepository clientRegistrationRepository;
+	private Map<OAuth2AuthorizedClientId, OAuth2AuthorizedClient> authorizedClients = new ConcurrentHashMap<>();
 
 	/**
 	 * Constructs an {@code InMemoryOAuth2AuthorizedClientService} using the provided parameters.
@@ -49,6 +49,15 @@ public InMemoryOAuth2AuthorizedClientService(ClientRegistrationRepository client
 		this.clientRegistrationRepository = clientRegistrationRepository;
 	}
 
+	/**
+	 * Sets the map of authorized clients to use.
+	 * @param authorizedClients the map of authorized clients
+	 */
+	public void setAuthorizedClients(Map<OAuth2AuthorizedClientId, OAuth2AuthorizedClient> authorizedClients) {
+		Assert.notNull(authorizedClients, "authorizedClients cannot be null");
+		this.authorizedClients = authorizedClients;
+	}
+
 	@Override
 	@SuppressWarnings("unchecked")
 	public <T extends OAuth2AuthorizedClient> T loadAuthorizedClient(String clientRegistrationId, String principalName) {

oauth2/oauth2-client/src/test/java/org/springframework/security/oauth2/client/InMemoryOAuth2AuthorizedClientServiceTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,7 +15,11 @@
  */
 package org.springframework.security.oauth2.client;
 
+import java.util.Collections;
+import java.util.Map;
+
 import org.junit.Test;
+
 import org.springframework.security.core.Authentication;
 import org.springframework.security.oauth2.client.registration.ClientRegistration;
 import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
@@ -24,13 +28,17 @@
 import org.springframework.security.oauth2.core.OAuth2AccessToken;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 /**
  * Tests for {@link InMemoryOAuth2AuthorizedClientService}.
  *
  * @author Joe Grandja
+ * @author Vedran Pavic
  */
 public class InMemoryOAuth2AuthorizedClientServiceTests {
 	private String principalName1 = "principal-1";
@@ -57,6 +65,30 @@ public void constructorWhenClientRegistrationRepositoryIsNullThenThrowIllegalArg
 		new InMemoryOAuth2AuthorizedClientService(null);
 	}
 
+	@Test
+	public void constructorWhenAuthorizedClientsIsNullThenIllegalArgumentException() {
+		assertThatExceptionOfType(IllegalArgumentException.class)
+				.isThrownBy(() -> this.authorizedClientService.setAuthorizedClients(null))
+				.withMessage("authorizedClients cannot be null");
+	}
+
+	@Test
+	public void constructorWhenAuthorizedClientsIsEmptyMapThenRepositoryUsingSuppliedAuthorizedClients() {
+		String registrationId = this.registration3.getRegistrationId();
+
+		Map<OAuth2AuthorizedClientId, OAuth2AuthorizedClient> authorizedClients = Collections.singletonMap(
+				OAuth2AuthorizedClientId.create(this.registration3, this.principalName1),
+				mock(OAuth2AuthorizedClient.class));
+		ClientRegistrationRepository clientRegistrationRepository = mock(ClientRegistrationRepository.class);
+		given(clientRegistrationRepository.findByRegistrationId(eq(registrationId))).willReturn(this.registration3);
+
+		InMemoryOAuth2AuthorizedClientService authorizedClientService = new InMemoryOAuth2AuthorizedClientService(
+				this.clientRegistrationRepository);
+		authorizedClientService.setAuthorizedClients(authorizedClients);
+		assertThat((OAuth2AuthorizedClient) authorizedClientService.loadAuthorizedClient(
+				registrationId, this.principalName1)).isNotNull();
+	}
+
 	@Test(expected = IllegalArgumentException.class)
 	public void loadAuthorizedClientWhenClientRegistrationIdIsNullThenThrowIllegalArgumentException() {
 		this.authorizedClientService.loadAuthorizedClient(null, this.principalName1);