diff --git a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md index 7034ec765efe..217622bdbf03 100644 --- a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md @@ -4,6 +4,12 @@ ### Features Added +- Added `getSessionState()` and `setSessionState(byte[])` to `ServiceBusReceivedMessageContext`, allowing + session state to be read and written directly from a `ServiceBusSessionProcessorClient` message handler + without opening a separate session receiver. The methods throw `IllegalStateException` when the message + was not received from a session-enabled entity, or when the session that delivered the message is no longer + held by the processor. ([#49207](https://github.com/Azure/azure-sdk-for-java/issues/49207)) + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java index f89e58154ef8..3438bf44497b 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java @@ -167,4 +167,37 @@ public void deadLetter(DeadLetterOptions options) { } receiverClient.deadLetter(receivedMessageContext.getMessage(), options).block(); } + + /** + * Gets the state of the session that the {@link #getMessage() message} in this context belongs to. This operation + * is only supported when the {@link ServiceBusProcessorClient} is consuming from a session-enabled entity. + * + * @return The session state, or {@code null} if no state is set for the session. + * @throws IllegalStateException if the message was not received from a session-enabled entity, or if the session + * that delivered the message is no longer held by this processor. + * @throws ServiceBusException if the session state could not be acquired. + */ + public byte[] getSessionState() { + if (sessionReceivers != null) { + return sessionReceivers.getSessionState(receivedMessageContext.getMessage()).block(); + } + return receiverClient.getSessionState(receivedMessageContext.getMessage()).block(); + } + + /** + * Sets the state of the session that the {@link #getMessage() message} in this context belongs to. This operation + * is only supported when the {@link ServiceBusProcessorClient} is consuming from a session-enabled entity. + * + * @param sessionState State to set on the session, or {@code null} to clear the session state. + * @throws IllegalStateException if the message was not received from a session-enabled entity, or if the session + * that delivered the message is no longer held by this processor. + * @throws ServiceBusException if the session state could not be set. + */ + public void setSessionState(byte[] sessionState) { + if (sessionReceivers != null) { + sessionReceivers.setSessionState(receivedMessageContext.getMessage(), sessionState).block(); + return; + } + receiverClient.setSessionState(receivedMessageContext.getMessage(), sessionState).block(); + } } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java index 0b3154ec33ef..38e71d22e60f 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java @@ -707,6 +707,19 @@ public Mono getSessionState() { return getSessionState(receiverOptions.getSessionId()); } + /** + * Gets the state of the session that delivered the given {@code message}. Unlike {@link #getSessionState()}, this + * targets the message's own session id rather than the receiver's bound session id, so it is correct for a session + * processor whose receiver is not bound to a single session. For a non-session receiver the underlying operation + * fails with an {@link IllegalStateException}. Package-private; invoked by {@link ServiceBusReceivedMessageContext}. + * + * @param message the message whose session state to read. + * @return the session state, or an empty Mono if no state is set for the session. + */ + Mono getSessionState(ServiceBusReceivedMessage message) { + return getSessionState(message.getSessionId()); + } + /** * Reads the next active message without changing the state of the receiver or the message source. The first call to * {@code peek()} fetches the first active message for this receiver. Each subsequent call fetches the subsequent @@ -1278,6 +1291,21 @@ public Mono setSessionState(byte[] sessionState) { return this.setSessionState(receiverOptions.getSessionId(), sessionState); } + /** + * Sets the state of the session that delivered the given {@code message}. Unlike {@link #setSessionState(byte[])}, + * this targets the message's own session id rather than the receiver's bound session id, so it is correct for a + * session processor whose receiver is not bound to a single session. For a non-session receiver the underlying + * operation fails with an {@link IllegalStateException}. Package-private; invoked by + * {@link ServiceBusReceivedMessageContext}. + * + * @param message the message whose session state to set. + * @param sessionState the state to set on the session, or {@code null} to clear it. + * @return a Mono that completes when the session state is set. + */ + Mono setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) { + return this.setSessionState(message.getSessionId(), sessionState); + } + /** * Starts a new service side transaction. The {@link ServiceBusTransactionContext transaction context} should be * passed to all operations that needs to be in this transaction. @@ -1843,6 +1871,9 @@ private Mono setSessionState(String sessionId, byte[] sessionState) { new IllegalStateException(String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "setSessionState"))); } else if (!isSessionEnabled) { return monoError(LOGGER, new IllegalStateException("Cannot set session state on a non-session receiver.")); + } else if (sessionId == null || sessionId.isEmpty()) { + return monoError(LOGGER, + new IllegalStateException("Cannot set session state because the session id is null or empty.")); } assert sessionManager != null; // guaranteed to be non-null when isSessionEnabled is true. final String linkName = sessionManager.getLinkName(sessionId); @@ -1860,12 +1891,15 @@ private Mono getSessionState(String sessionId) { new IllegalStateException(String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "getSessionState"))); } else if (!isSessionEnabled) { return monoError(LOGGER, new IllegalStateException("Cannot get session state on a non-session receiver.")); + } else if (sessionId == null || sessionId.isEmpty()) { + return monoError(LOGGER, + new IllegalStateException("Cannot get session state because the session id is null or empty.")); } assert sessionManager != null; // guaranteed to be non-null when isSessionEnabled is true. final String linkName = sessionManager.getLinkName(sessionId); return tracer - .traceMono("ServiceBus.setSessionState", + .traceMono("ServiceBus.getSessionState", connectionProcessor.flatMap(connection -> connection.getManagementNode(entityPath, entityType)) .flatMap(channel -> channel.getSessionState(sessionId, linkName))) .onErrorMap((err) -> mapError(err, ServiceBusErrorSource.RECEIVE)); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java index ee9a0a9e0399..ee3e8dc3a0ac 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java @@ -285,6 +285,16 @@ ServiceBusReceiveLink getLink() { return link; } + /** + * Gets the Mono that resolves the management node used for session-scoped operations such as session state + * get/set and session lock renewal. + * + * @return the Mono to obtain the management node for this session. + */ + Mono getSessionManagement() { + return sessionManagement; + } + /** * Begin the recurring lock renewal for the session. * diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java index 42cb50c79f8b..1b35f4a54337 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java @@ -6,6 +6,7 @@ import com.azure.core.amqp.implementation.AmqpReceiveLink; import com.azure.core.util.logging.ClientLogger; import com.azure.core.util.logging.LoggingEventBuilder; +import com.azure.messaging.servicebus.implementation.ServiceBusManagementNode; import com.azure.messaging.servicebus.implementation.instrumentation.ServiceBusTracer; import org.apache.qpid.proton.amqp.transport.DeliveryState; import org.apache.qpid.proton.message.Message; @@ -24,8 +25,10 @@ final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink { private final ClientLogger logger; + private final ServiceBusTracer tracer; private final String sessionId; private final AmqpReceiveLink sessionLink; + private final Mono sessionManagement; private final boolean hasIdleTimeout; private final Sinks.Many nextItemIdleTimeoutSink = Sinks.many().multicast().onBackpressureBuffer(); private final Sinks.Empty terminateEndpointStatesSink = Sinks.empty(); @@ -34,8 +37,10 @@ final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink { ServiceBusSessionReactorReceiver(ClientLogger logger, ServiceBusTracer tracer, ServiceBusSessionAcquirer.Session session, Duration sessionIdleTimeout, Duration maxSessionLockRenew) { this.logger = logger; + this.tracer = tracer; this.sessionId = session.getId(); this.sessionLink = session.getLink(); + this.sessionManagement = session.getSessionManagement(); this.hasIdleTimeout = sessionIdleTimeout != null; if (hasIdleTimeout) { this.disposables @@ -53,6 +58,41 @@ public String getSessionId() { return sessionId; } + /** + * Gets the state of the session this receiver is streaming messages from. + * + * @return A Mono that completes with the session state, or an empty Mono if there is no state set for the session. + */ + Mono getSessionState() { + return tracer + .traceMono("ServiceBus.getSessionState", + sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName()))) + .onErrorMap(ServiceBusSessionReactorReceiver::mapSessionStateError); + } + + /** + * Sets the state of the session this receiver is streaming messages from. + * + * @param sessionState State to set on the session, or {@code null} to clear the session state. + * @return A Mono that completes when the session state is set. + */ + Mono setSessionState(byte[] sessionState) { + return tracer + .traceMono("ServiceBus.setSessionState", + sessionManagement.flatMap(mgmt -> mgmt.setSessionState(sessionId, sessionState, getLinkName()))) + .onErrorMap(ServiceBusSessionReactorReceiver::mapSessionStateError); + } + + // Mirror ServiceBusReceiverAsyncClient's session-state error handling: surface a ServiceBusException to the + // ServiceBusReceivedMessageContext caller (its Javadoc documents ServiceBusException), rather than a raw + // AmqpException from the management link. + private static Throwable mapSessionStateError(Throwable throwable) { + if (throwable instanceof ServiceBusException) { + return throwable; + } + return new ServiceBusException(throwable, ServiceBusErrorSource.RECEIVE); + } + @Override public String getHostname() { return sessionLink.getHostname(); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java index 091f81ebee2c..80cd0aaa9af6 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java @@ -789,8 +789,10 @@ static final class SessionReceiversTracker { private final ConcurrentHashMap receivers; private final ServiceBusReceiverInstrumentation instrumentation; - private SessionReceiversTracker(ClientLogger logger, int size, String fullyQualifiedNamespace, - String entityPath, ServiceBusReceiveMode receiveMode, ServiceBusReceiverInstrumentation instrumentation) { + // Visible for testing (package-private) so SessionReceiversTrackerTest can construct the tracker and + // exercise the session-state paths directly; in production it is only constructed by SessionsMessagePump. + SessionReceiversTracker(ClientLogger logger, int size, String fullyQualifiedNamespace, String entityPath, + ServiceBusReceiveMode receiveMode, ServiceBusReceiverInstrumentation instrumentation) { this.logger = logger; this.fullyQualifiedNamespace = fullyQualifiedNamespace; this.entityPath = entityPath; @@ -799,7 +801,8 @@ private SessionReceiversTracker(ClientLogger logger, int size, String fullyQuali this.instrumentation = instrumentation; } - private void track(ServiceBusSessionReactorReceiver receiver) { + // Visible for testing (package-private) so SessionReceiversTrackerTest can populate the receivers map. + void track(ServiceBusSessionReactorReceiver receiver) { receivers.put(receiver.getSessionId().toLowerCase(Locale.ROOT), receiver); } @@ -872,6 +875,43 @@ Mono defer(ServiceBusReceivedMessage message, DeferOptions options) { options.getTransactionContext()); } + Mono getSessionState(ServiceBusReceivedMessage message) { + final String sessionId = message.getSessionId(); + if (sessionId == null || sessionId.isEmpty()) { + return notSessionMessageError("getSessionState"); + } + final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT)); + if (receiver == null) { + return sessionNotActiveError(sessionId, "getSessionState"); + } + return receiver.getSessionState(); + } + + Mono setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) { + final String sessionId = message.getSessionId(); + if (sessionId == null || sessionId.isEmpty()) { + return notSessionMessageError("setSessionState"); + } + final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT)); + if (receiver == null) { + return sessionNotActiveError(sessionId, "setSessionState"); + } + return receiver.setSessionState(sessionState); + } + + private Mono notSessionMessageError(String operation) { + return monoError(logger, new IllegalStateException(String.format( + "Cannot perform '%s'; the message was not received from a session-enabled entity.", operation))); + } + + private Mono sessionNotActiveError(String sessionId, String operation) { + final String m = String.format( + "Cannot perform '%s'; the session '%s' is no longer active for this processor. Session state " + + "operations are only valid while the session that delivered the message is still held.", + operation, sessionId); + return monoError(logger, new IllegalStateException(m)); + } + private Mono updateDisposition(ServiceBusReceivedMessage message, DispositionStatus dispositionStatus, Map propertiesToModify, String deadLetterReason, String deadLetterDescription, ServiceBusTransactionContext transactionContext) { diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java new file mode 100644 index 000000000000..857063280a98 --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java @@ -0,0 +1,174 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.BinaryData; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import reactor.core.publisher.Mono; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link ServiceBusReceivedMessageContext}, focused on the session-state operations exposed to a + * {@link ServiceBusProcessorClient} message handler. Covers both delegation branches: the receiver-client path + * (non-session and V1 processors) and the V2 session-processor path ({@code SessionsMessagePump}). + */ +class ServiceBusReceivedMessageContextTest { + private static final byte[] STATE = new byte[] { 10, 111, 23 }; + + @Mock + private ServiceBusReceiverAsyncClient receiverClient; + + private AutoCloseable mocksCloseable; + private ServiceBusReceivedMessage message; + private ServiceBusMessageContext messageContext; + + @BeforeEach + void setup() { + mocksCloseable = MockitoAnnotations.openMocks(this); + message = new ServiceBusReceivedMessage(BinaryData.fromString("some-data")); + messageContext = new ServiceBusMessageContext(message); + } + + @AfterEach + void teardown() throws Exception { + if (mocksCloseable != null) { + mocksCloseable.close(); + } + } + + // --- Receiver-client path: non-session processor and V1 processor --- + + @Test + void getSessionStateDelegatesToReceiver() { + when(receiverClient.getSessionState(message)).thenReturn(Mono.just(STATE)); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertArrayEquals(STATE, context.getSessionState()); + verify(receiverClient).getSessionState(message); + } + + @Test + void getSessionStateReturnsNullWhenNoState() { + when(receiverClient.getSessionState(message)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertNull(context.getSessionState()); + } + + @Test + void setSessionStateDelegatesToReceiver() { + when(receiverClient.setSessionState(message, STATE)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + context.setSessionState(STATE); + + verify(receiverClient).setSessionState(message, STATE); + } + + /** + * Asymmetry test: on a non-session entity the underlying receiver rejects session-state reads, and the context + * must surface that {@link IllegalStateException} rather than silently succeeding. + */ + @Test + void getSessionStateOnNonSessionEntityThrows() { + when(receiverClient.getSessionState(message)) + .thenReturn(Mono.error(new IllegalStateException("Cannot get session state on a non-session receiver."))); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertThrows(IllegalStateException.class, context::getSessionState); + } + + /** + * Asymmetry test: on a non-session entity the underlying receiver rejects session-state writes, and the context + * must surface that {@link IllegalStateException} rather than silently succeeding. + */ + @Test + void setSessionStateOnNonSessionEntityThrows() { + when(receiverClient.setSessionState(message, STATE)) + .thenReturn(Mono.error(new IllegalStateException("Cannot set session state on a non-session receiver."))); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertThrows(IllegalStateException.class, () -> context.setSessionState(STATE)); + } + + // --- Session-processor path: V2 SessionsMessagePump (the #49207 scenario) --- + + @Test + void getSessionStateDelegatesToSessionTracker() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.getSessionState(message)).thenReturn(Mono.just(STATE)); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertArrayEquals(STATE, context.getSessionState()); + verify(tracker).getSessionState(message); + } + + @Test + void getSessionStateViaSessionTrackerReturnsNullWhenNoState() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.getSessionState(message)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertNull(context.getSessionState()); + } + + @Test + void setSessionStateDelegatesToSessionTracker() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.setSessionState(message, STATE)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + context.setSessionState(STATE); + + verify(tracker).setSessionState(message, STATE); + } + + /** + * Asymmetry test: when the session that delivered the message is no longer held by the processor, the tracker + * errors and the context must surface that {@link IllegalStateException} for a read. + */ + @Test + void getSessionStateThrowsWhenSessionNoLongerActive() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.getSessionState(message)) + .thenReturn(Mono.error(new IllegalStateException("the session is no longer active for this processor."))); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertThrows(IllegalStateException.class, context::getSessionState); + } + + /** + * Asymmetry test: when the session that delivered the message is no longer held by the processor, the tracker + * errors and the context must surface that {@link IllegalStateException} for a write. + */ + @Test + void setSessionStateThrowsWhenSessionNoLongerActive() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.setSessionState(message, STATE)) + .thenReturn(Mono.error(new IllegalStateException("the session is no longer active for this processor."))); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertThrows(IllegalStateException.class, () -> context.setSessionState(STATE)); + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java new file mode 100644 index 000000000000..104558c9f25b --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.BinaryData; +import com.azure.core.util.logging.ClientLogger; +import com.azure.messaging.servicebus.models.ServiceBusReceiveMode; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import java.util.Arrays; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link SessionsMessagePump.SessionReceiversTracker} session-state routing (#49207): looking up the + * session receiver by the message's session id, erroring when the session is not (or no longer) tracked, and + * delegating to the tracked {@link ServiceBusSessionReactorReceiver} otherwise. + */ +class SessionReceiversTrackerTest { + private static final byte[] STATE = new byte[] { 5, 6, 7 }; + + private static SessionsMessagePump.SessionReceiversTracker newTracker() { + return new SessionsMessagePump.SessionReceiversTracker(new ClientLogger(SessionReceiversTrackerTest.class), 4, + "fqdn.servicebus.windows.net", "entity", ServiceBusReceiveMode.PEEK_LOCK, null); + } + + private static ServiceBusReceivedMessage messageForSession(String sessionId) { + final ServiceBusReceivedMessage message = new ServiceBusReceivedMessage(BinaryData.fromString("data")); + if (sessionId != null) { + message.setSessionId(sessionId); + } + return message; + } + + @Test + void getSessionStateErrorsWhenMessageHasNoSession() { + StepVerifier.create(newTracker().getSessionState(messageForSession(null))) + .expectError(IllegalStateException.class) + .verify(); + } + + @Test + void getSessionStateErrorsWhenMessageHasEmptySession() { + StepVerifier.create(newTracker().getSessionState(messageForSession(""))) + .expectErrorMatches(e -> e instanceof IllegalStateException + && e.getMessage().contains("was not received from a session-enabled entity")) + .verify(); + } + + @Test + void setSessionStateErrorsWhenMessageHasEmptySession() { + StepVerifier.create(newTracker().setSessionState(messageForSession(""), STATE)) + .expectErrorMatches(e -> e instanceof IllegalStateException + && e.getMessage().contains("was not received from a session-enabled entity")) + .verify(); + } + + @Test + void getSessionStateErrorsWhenSessionNotTracked() { + StepVerifier.create(newTracker().getSessionState(messageForSession("absent-session"))) + .expectError(IllegalStateException.class) + .verify(); + } + + @Test + void setSessionStateErrorsWhenSessionNotTracked() { + StepVerifier.create(newTracker().setSessionState(messageForSession("absent-session"), STATE)) + .expectError(IllegalStateException.class) + .verify(); + } + + @Test + void getSessionStateDelegatesToTrackedReceiver() { + final SessionsMessagePump.SessionReceiversTracker tracker = newTracker(); + final ServiceBusSessionReactorReceiver receiver = mock(ServiceBusSessionReactorReceiver.class); + when(receiver.getSessionId()).thenReturn("s1"); + when(receiver.getSessionState()).thenReturn(Mono.just(STATE)); + tracker.track(receiver); + + StepVerifier.create(tracker.getSessionState(messageForSession("s1"))) + .expectNextMatches(actual -> Arrays.equals(STATE, actual)) + .verifyComplete(); + verify(receiver).getSessionState(); + } + + @Test + void setSessionStateDelegatesToTrackedReceiver() { + final SessionsMessagePump.SessionReceiversTracker tracker = newTracker(); + final ServiceBusSessionReactorReceiver receiver = mock(ServiceBusSessionReactorReceiver.class); + when(receiver.getSessionId()).thenReturn("s1"); + when(receiver.setSessionState(STATE)).thenReturn(Mono.empty()); + tracker.track(receiver); + + StepVerifier.create(tracker.setSessionState(messageForSession("s1"), STATE)).verifyComplete(); + verify(receiver).setSessionState(STATE); + } + + /** + * Session-id lookup is case-insensitive: a receiver tracked under one casing is found by a message whose session + * id differs only in case (mirrors how {@code track()} and the settlement path key the receivers map). + */ + @Test + void getSessionStateLookupIsCaseInsensitive() { + final SessionsMessagePump.SessionReceiversTracker tracker = newTracker(); + final ServiceBusSessionReactorReceiver receiver = mock(ServiceBusSessionReactorReceiver.class); + when(receiver.getSessionId()).thenReturn("Session-A"); + when(receiver.getSessionState()).thenReturn(Mono.just(STATE)); + tracker.track(receiver); + + StepVerifier.create(tracker.getSessionState(messageForSession("SESSION-A"))) + .expectNextMatches(actual -> Arrays.equals(STATE, actual)) + .verifyComplete(); + } +}