Improve WebSession related tests

Add missing DefaultWebSessionManagerTests .block(). Previously
session.save() was invoked, but we did not ensure it was completed.
This commit makes it block on session.save()

Fix existingSessionIsExpired. This test is actually broken and is
testing a new session is created because the session id returned
by the idResolver does not match the existing WebSession.
This commit ensures that the id of the WebSession found by idResolver
matches the existing WebSession.

DefaultWebSessionManagerTests use Mockito. To ensure we test with
independence from InMemoryWebSessionStore we use Mockito for the
DefaultWebessionManager collaborators.

Add test for response.setComplete(). We want to ensure that when the
response is completed, it saves the WebSession and writes it to the
response using idResolver

spring-web/src/test/java/org/springframework/web/server/session/DefaultWebSessionManagerTests.java

@@ -19,17 +19,17 @@ import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.time.ZoneId;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
 import reactor.core.publisher.Mono;
 
 import org.springframework.http.codec.ServerCodecConfigurer;
-import org.springframework.lang.Nullable;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
 import org.springframework.util.IdGenerator;
@@ -43,13 +43,18 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNotSame;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertSame;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 /**
  * Unit tests for {@link DefaultWebSessionManager}.
  * @author Rossen Stoyanchev
+ * @author Rob Winch
  */
+@RunWith(MockitoJUnitRunner.class)
 public class DefaultWebSessionManagerTests {
 
 	private static final Clock CLOCK = Clock.system(ZoneId.of("GMT"));
@@ -59,16 +64,19 @@ public class DefaultWebSessionManagerTests {
 
 	private DefaultWebSessionManager manager;
 
-	private TestWebSessionIdResolver idResolver;
-
 	private ServerWebExchange exchange;
 
+	@Mock
+	private WebSessionIdResolver idResolver;
+
+	@Mock
+	private WebSessionStore store;
 
 	@Before
 	public void setUp() throws Exception {
 		this.manager = new DefaultWebSessionManager();
-		this.idResolver = new TestWebSessionIdResolver();
 		this.manager.setSessionIdResolver(this.idResolver);
+		this.manager.setSessionStore(this.store);
 
 		MockServerHttpRequest request = MockServerHttpRequest.get("/path").build();
 		MockServerHttpResponse response = new MockServerHttpResponse();
@@ -79,45 +87,60 @@ public class DefaultWebSessionManagerTests {
 
 	@Test
 	public void getSessionWithoutStarting() throws Exception {
-		this.idResolver.setIdsToResolve(Collections.emptyList());
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Collections.emptyList());
 		WebSession session = this.manager.getSession(this.exchange).block();
-		session.save();
+		session.save().block();
 
 		assertFalse(session.isStarted());
 		assertFalse(session.isExpired());
-		assertNull(this.idResolver.getSavedId());
-		assertNull(this.manager.getSessionStore().retrieveSession(session.getId()).block());
+		verify(this.store, never()).storeSession(any());
+		verify(this.idResolver, never()).setSessionId(any(), any());
 	}
 
 	@Test
 	public void startSessionExplicitly() throws Exception {
-		this.idResolver.setIdsToResolve(Collections.emptyList());
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Collections.emptyList());
+		when(this.store.storeSession(any())).thenReturn(Mono.empty());
 		WebSession session = this.manager.getSession(this.exchange).block();
 		session.start();
-		session.save();
+		session.save().block();
 
 		String id = session.getId();
-		assertNotNull(this.idResolver.getSavedId());
-		assertEquals(id, this.idResolver.getSavedId());
-		assertSame(session, this.manager.getSessionStore().retrieveSession(id).block());
+		verify(this.store).storeSession(any());
+		verify(this.idResolver).setSessionId(any(), eq(id));
 	}
 
 	@Test
 	public void startSessionImplicitly() throws Exception {
-		this.idResolver.setIdsToResolve(Collections.emptyList());
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Collections.emptyList());
+		when(this.store.storeSession(any())).thenReturn(Mono.empty());
 		WebSession session = this.manager.getSession(this.exchange).block();
 		session.getAttributes().put("foo", "bar");
-		session.save();
+		session.save().block();
 
-		assertNotNull(this.idResolver.getSavedId());
+		verify(this.idResolver).setSessionId(any(), any());
+		verify(this.store).storeSession(any());
+	}
+
+	@Test
+	public void exchangeWhenResponseSetCompleteThenSavesAndSetsId() throws Exception {
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Collections.emptyList());
+		when(this.store.storeSession(any())).thenReturn(Mono.empty());
+		WebSession session = this.manager.getSession(this.exchange).block();
+		String id = session.getId();
+		session.getAttributes().put("foo", "bar");
+		this.exchange.getResponse().setComplete().block();
+
+		verify(this.idResolver).setSessionId(any(), eq(id));
+		verify(this.store).storeSession(any());
 	}
 
 	@Test
 	public void existingSession() throws Exception {
 		DefaultWebSession existing = createDefaultWebSession();
 		String id = existing.getId();
-		this.manager.getSessionStore().storeSession(existing);
-		this.idResolver.setIdsToResolve(Collections.singletonList(id));
+		when(this.store.retrieveSession(id)).thenReturn(Mono.just(existing));
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Collections.singletonList(id));
 
 		WebSession actual = this.manager.getSession(this.exchange).block();
 		assertNotNull(actual);
@@ -130,19 +153,23 @@ public class DefaultWebSessionManagerTests {
 		existing.start();
 		Instant lastAccessTime = Instant.now(CLOCK).minus(Duration.ofMinutes(31));
 		existing = new DefaultWebSession(existing, lastAccessTime, s -> Mono.empty());
-		this.manager.getSessionStore().storeSession(existing);
-		this.idResolver.setIdsToResolve(Collections.singletonList("1"));
+		when(this.store.retrieveSession(existing.getId())).thenReturn(Mono.just(existing));
+		when(this.store.removeSession(existing.getId())).thenReturn(Mono.empty());
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Collections.singletonList(existing.getId()));
 
 		WebSession actual = this.manager.getSession(this.exchange).block();
 		assertNotSame(existing, actual);
+		verify(this.store).removeSession(existing.getId());
+		verify(this.idResolver).expireSession(any());
 	}
 
 	@Test
 	public void multipleSessionIds() throws Exception {
 		DefaultWebSession existing = createDefaultWebSession();
 		String id = existing.getId();
-		this.manager.getSessionStore().storeSession(existing);
-		this.idResolver.setIdsToResolve(Arrays.asList("neither-this", "nor-that", id));
+		when(this.store.retrieveSession(any())).thenReturn(Mono.empty());
+		when(this.store.retrieveSession(id)).thenReturn(Mono.just(existing));
+		when(this.idResolver.resolveSessionIds(this.exchange)).thenReturn(Arrays.asList("neither-this", "nor-that", id));
 
 		WebSession actual = this.manager.getSession(this.exchange).block();
 		assertNotNull(actual);
@@ -152,39 +179,4 @@ public class DefaultWebSessionManagerTests {
 	private DefaultWebSession createDefaultWebSession() {
 		return new DefaultWebSession(idGenerator, CLOCK, (s, session) -> Mono.empty(), s -> Mono.empty());
 	}
-
-
-	private static class TestWebSessionIdResolver implements WebSessionIdResolver {
-
-		private List<String> idsToResolve = new ArrayList<>();
-
-		@Nullable
-		private String id = null;
-
-
-		public void setIdsToResolve(List<String> idsToResolve) {
-			this.idsToResolve = idsToResolve;
-		}
-
-		@Nullable
-		public String getSavedId() {
-			return this.id;
-		}
-
-		@Override
-		public List<String> resolveSessionIds(ServerWebExchange exchange) {
-			return this.idsToResolve;
-		}
-
-		@Override
-		public void setSessionId(ServerWebExchange exchange, String sessionId) {
-			this.id = sessionId;
-		}
-
-		@Override
-		public void expireSession(ServerWebExchange exchange) {
-			this.id = null;
-		}
-	}
-
 }