Refactor WebSession#getAttribute options

Issue: SPR-15718

spring-test/src/test/java/org/springframework/test/web/reactive/server/MockServerIntegrationTests.java

@@ -43,9 +43,9 @@ public class MockServerIntegrationTests {
 			}
 			else {
 				return exchange.getSession()
-						.map(session -> session.getAttribute("foo").orElse("none"))
+						.map(session -> session.getAttributeOrDefault("foo", "none"))
 						.flatMap(value -> {
-							byte[] bytes = value.toString().getBytes(UTF_8);
+							byte[] bytes = value.getBytes(UTF_8);
 							DataBuffer buffer = new DefaultDataBufferFactory().wrap(bytes);
 							return exchange.getResponse().writeWith(Mono.just(buffer));
 						});

spring-web/src/main/java/org/springframework/web/server/WebSession.java

@@ -19,10 +19,12 @@ package org.springframework.web.server;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Map;
-import java.util.Optional;
 
 import reactor.core.publisher.Mono;
 
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
 /**
  * Main contract for using a server-side session that provides access to session
  * attributes across HTTP requests.
@@ -48,12 +50,42 @@ public interface WebSession {
 	Map<String, Object> getAttributes();
 
 	/**
-	 * Return the attribute value if present.
+	 * Return the session attribute value if present.
+	 * @param name the attribute name
+	 * @param <T> the attribute type
+	 * @return the attribute value
+	 */
+	@SuppressWarnings("unchecked")
+	@Nullable
+	default <T> T getAttribute(String name) {
+		return (T) getAttributes().get(name);
+	}
+
+	/**
+	 * Return the session attribute value or if not present raise an
+	 * {@link IllegalArgumentException}.
+	 * @param name the attribute name
+	 * @param <T> the attribute type
+	 * @return the attribute value
+	 */
+	@SuppressWarnings("unchecked")
+	default <T> T getRequiredAttribute(String name) {
+		T value = getAttribute(name);
+		Assert.notNull(value, "Required attribute '" + name + "' is missing.");
+		return value;
+	}
+
+	/**
+	 * Return the session attribute value, or a default, fallback value.
 	 * @param name the attribute name
+	 * @param defaultValue a default value to return instead
 	 * @param <T> the attribute type
 	 * @return the attribute value
 	 */
-	<T> Optional<T> getAttribute(String name);
+	@SuppressWarnings("unchecked")
+	default <T> T getAttributeOrDefault(String name, T defaultValue) {
+		return (T) getAttributes().getOrDefault(name, defaultValue);
+	}
 
 	/**
 	 * Force the creation of a session causing the session id to be sent when

spring-web/src/main/java/org/springframework/web/server/session/DefaultWebSession.java

@@ -20,7 +20,6 @@ import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
@@ -108,11 +107,6 @@ public class DefaultWebSession implements ConfigurableWebSession, Serializable {
 		return this.attributes;
 	}
 
-	@Override @SuppressWarnings("unchecked")
-	public <T> Optional<T> getAttribute(String name) {
-		return Optional.ofNullable((T) this.attributes.get(name));
-	}
-
 	@Override
 	public Instant getCreationTime() {
 		return this.creationTime;

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/SessionAttributeMethodArgumentResolver.java

@@ -58,9 +58,8 @@ public class SessionAttributeMethodArgumentResolver extends AbstractNamedValueAr
 	@Override
 	protected Mono<Object> resolveName(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		return exchange.getSession()
-				.map(session -> session.getAttribute(name))
-				.filter(Optional::isPresent)
-				.map(Optional::get);
+				.filter(session -> session.getAttribute(name) != null)
+				.map(session -> session.getAttribute(name));
 	}
 
 	@Override

spring-webflux/src/test/java/org/springframework/web/reactive/result/method/annotation/SessionAttributeMethodArgumentResolverTests.java

@@ -47,7 +47,6 @@ import org.springframework.web.server.i18n.AcceptHeaderLocaleContextResolver;
 import org.springframework.web.server.session.MockWebSessionManager;
 import org.springframework.web.server.session.WebSessionManager;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
@@ -103,7 +102,7 @@ public class SessionAttributeMethodArgumentResolverTests {
 		StepVerifier.create(mono).expectError(ServerWebInputException.class).verify();
 
 		Foo foo = new Foo();
-		when(this.session.getAttribute("foo")).thenReturn(Optional.of(foo));
+		when(this.session.getAttribute("foo")).thenReturn(foo);
 		mono = this.resolver.resolveArgument(param, new BindingContext(), this.exchange);
 		assertSame(foo, mono.block());
 	}
@@ -112,7 +111,7 @@ public class SessionAttributeMethodArgumentResolverTests {
 	public void resolveWithName() throws Exception {
 		MethodParameter param = initMethodParameter(1);
 		Foo foo = new Foo();
-		when(this.session.getAttribute("specialFoo")).thenReturn(Optional.of(foo));
+		when(this.session.getAttribute("specialFoo")).thenReturn(foo);
 		Mono<Object> mono = this.resolver.resolveArgument(param, new BindingContext(), this.exchange);
 		assertSame(foo, mono.block());
 	}
@@ -124,32 +123,32 @@ public class SessionAttributeMethodArgumentResolverTests {
 		assertNull(mono.block());
 
 		Foo foo = new Foo();
-		when(this.session.getAttribute("foo")).thenReturn(Optional.of(foo));
+		when(this.session.getAttribute("foo")).thenReturn(foo);
 		mono = this.resolver.resolveArgument(param, new BindingContext(), this.exchange);
 		assertSame(foo, mono.block());
 	}
 
+	@SuppressWarnings("unchecked")
 	@Test
 	public void resolveOptional() throws Exception {
 		MethodParameter param = initMethodParameter(3);
-		Mono<Object> mono = this.resolver.resolveArgument(param, new BindingContext(), this.exchange);
-		assertNotNull(mono.block());
-		assertEquals(Optional.class, mono.block().getClass());
-		assertFalse(((Optional<?>) mono.block()).isPresent());
+		Optional<Object> actual = (Optional<Object>) this.resolver
+				.resolveArgument(param, new BindingContext(), this.exchange).block();
+
+		assertNotNull(actual);
+		assertFalse(actual.isPresent());
 
 		ConfigurableWebBindingInitializer initializer = new ConfigurableWebBindingInitializer();
 		initializer.setConversionService(new DefaultFormattingConversionService());
 		BindingContext bindingContext = new BindingContext(initializer);
 
 		Foo foo = new Foo();
-		when(this.session.getAttribute("foo")).thenReturn(Optional.of(foo));
-		mono = this.resolver.resolveArgument(param, bindingContext, this.exchange);
-
-		assertNotNull(mono.block());
-		assertEquals(Optional.class, mono.block().getClass());
-		Optional<?> optional = (Optional<?>) mono.block();
-		assertTrue(optional.isPresent());
-		assertSame(foo, optional.get());
+		when(this.session.getAttribute("foo")).thenReturn(foo);
+		actual = (Optional<Object>) this.resolver.resolveArgument(param, bindingContext, this.exchange).block();
+
+		assertNotNull(actual);
+		assertTrue(actual.isPresent());
+		assertSame(foo, actual.get());
 	}