tsai
/
thingsboard
mirror of https://github.com/thingsboard/thingsboard.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

added tests for DefaultWebSocketService

pull/15313/head
dashevchenko 6 months ago
parent
a75c008f50
commit
bf307a41bd
2 changed files with 111 additions and 2 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 6
      application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java
  2. 107
      application/src/test/java/org/thingsboard/server/service/ws/DefaultWebSocketServiceTest.java

6
application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java
View File

@ -315,7 +315,7 @@ public class DefaultWebSocketService implements WebSocketService {
}
}
private void processSessionClose(WebSocketSessionRef sessionRef) {
void processSessionClose(WebSocketSessionRef sessionRef) {
var tenantProfileConfiguration = getTenantProfileConfiguration(sessionRef);
if (tenantProfileConfiguration != null) {
String sessionId = "[" + sessionRef.getSessionId() + "]";
@ -403,7 +403,9 @@ public class DefaultWebSocketService implements WebSocketService {
if (tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser() > 0 && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) {
Set<String> publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet());
synchronized (publicUserSessions) {
if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser()) {
if (cmd.isUnsubscribe()) {
publicUserSessions.remove(subId);
} else if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser()) {
publicUserSessions.add(subId);
} else {
log.info("[{}][{}][{}] Failed to start subscription. Max public user subscriptions limit reached"

107
application/src/test/java/org/thingsboard/server/service/ws/DefaultWebSocketServiceTest.java
View File

@ -146,6 +146,113 @@ class DefaultWebSocketServiceTest {
assertThat(subs.iterator().next()).isEqualTo("[" + sessionId + "]:[" + cmdId + "]");
}
@Test
void processSubscription_unsubscribe_removesEntryFromPublicUserSubscriptionsMap() {
int maxPublicSubscriptions = 5;
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID());
TenantProfile profile = new TenantProfile();
profile.createDefaultTenantProfileData();
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions);
willReturn(profile).given(tenantProfileCache).get(tenantId);
String sessionId = "session-1";
int cmdId = 1;
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId);
service.processSubscription(sessionRef, subscriptionCmd(cmdId));
@SuppressWarnings("unchecked")
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap =
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap");
assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(1);
AttributesSubscriptionCmd unsubCmd = subscriptionCmd(cmdId);
unsubCmd.setUnsubscribe(true);
service.processSubscription(sessionRef, unsubCmd);
assertThat(publicUserSubscriptionsMap.get(tenantId)).isEmpty();
}
@Test
void processSubscription_unsubscribe_freesSlotForNewSubscription() {
int maxPublicSubscriptions = 1;
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID());
TenantProfile profile = new TenantProfile();
profile.createDefaultTenantProfileData();
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions);
willReturn(profile).given(tenantProfileCache).get(tenantId);
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, "session-1");
service.processSubscription(sessionRef, subscriptionCmd(1));
// slot is full — second subscription on same session should be rejected
assertThat(service.processSubscription(sessionRef, subscriptionCmd(2))).isFalse();
// unsubscribe cmd 1 to free the slot
AttributesSubscriptionCmd unsubCmd = subscriptionCmd(1);
unsubCmd.setUnsubscribe(true);
service.processSubscription(sessionRef, unsubCmd);
// now a new subscription should succeed
assertThat(service.processSubscription(sessionRef, subscriptionCmd(3)))
.as("new subscription should succeed after unsubscribe freed the slot")
.isTrue();
}
@Test
void processSessionClose_removesAllSessionSubscriptionsFromPublicUserSubscriptionsMap() {
int maxPublicSubscriptions = 10;
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID());
TenantProfile profile = new TenantProfile();
profile.createDefaultTenantProfileData();
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions);
willReturn(profile).given(tenantProfileCache).get(tenantId);
String sessionId = "closing-session";
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId);
service.processSubscription(sessionRef, subscriptionCmd(1));
service.processSubscription(sessionRef, subscriptionCmd(2));
service.processSubscription(sessionRef, subscriptionCmd(3));
@SuppressWarnings("unchecked")
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap =
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap");
assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(3);
service.processSessionClose(sessionRef);
assertThat(publicUserSubscriptionsMap.get(tenantId)).isEmpty();
}
@Test
void processSessionClose_onlyRemovesClosedSessionSubscriptions() {
int maxPublicSubscriptions = 10;
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID());
TenantProfile profile = new TenantProfile();
profile.createDefaultTenantProfileData();
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions);
willReturn(profile).given(tenantProfileCache).get(tenantId);
WebSocketSessionRef session1 = mockPublicSessionRef(tenantId, "session-1");
WebSocketSessionRef session2 = mockPublicSessionRef(tenantId, "session-2");
service.processSubscription(session1, subscriptionCmd(1));
service.processSubscription(session1, subscriptionCmd(2));
service.processSubscription(session2, subscriptionCmd(1));
@SuppressWarnings("unchecked")
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap =
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap");
assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(3);
service.processSessionClose(session1);
Set<String> remaining = publicUserSubscriptionsMap.get(tenantId);
assertThat(remaining).hasSize(1);
assertThat(remaining).allMatch(subId -> subId.startsWith("[session-2]"));
}
private WebSocketSessionRef mockPublicSessionRef(TenantId tenantId, String sessionId) {
CustomerId customerId = new CustomerId(UUID.randomUUID());
SecurityUser securityUser = mock(SecurityUser.class);

Write Preview
Loading…
Cancel
Save