From 873ad00792730a29f87325b74c86dd51cc3dc5b9 Mon Sep 17 00:00:00 2001
From: YevhenBondarenko <ybondarenko@thingsboard.io>
Date: Wed, 2 Oct 2024 16:30:17 +0200
Subject: [PATCH] minor refactoring and test improvements due to comments

---
 .../subscription/TbEntityLocalSubsInfo.java   |  6 ++--
 .../TbEntityLocalSubsInfoTest.java            | 33 ++++++++++++-------
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/application/src/main/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfo.java b/application/src/main/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfo.java
index 12786c65ee..ee20843538 100644
--- a/application/src/main/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfo.java
+++ b/application/src/main/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfo.java
@@ -135,7 +135,7 @@ public class TbEntityLocalSubsInfo {
         }
         TbSubscriptionType type = sub.getType();
         TbSubscriptionsInfo newState = state.copy();
-        updateNewState(newState, type);
+        clearState(newState, type);
         return updateState(Set.of(type), newState);
     }
 
@@ -155,14 +155,14 @@ public class TbEntityLocalSubsInfo {
                 continue;
             }
 
-            updateNewState(newState, type);
+            clearState(newState, type);
             changedTypes.add(type);
         }
 
         return updateState(changedTypes, newState);
     }
 
-    private void updateNewState(TbSubscriptionsInfo state, TbSubscriptionType type) {
+    private void clearState(TbSubscriptionsInfo state, TbSubscriptionType type) {
         switch (type) {
             case NOTIFICATIONS:
             case NOTIFICATIONS_COUNT:
diff --git a/application/src/test/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfoTest.java b/application/src/test/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfoTest.java
index 88f29308ab..e9b95ec832 100644
--- a/application/src/test/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfoTest.java
+++ b/application/src/test/java/org/thingsboard/server/service/subscription/TbEntityLocalSubsInfoTest.java
@@ -21,7 +21,6 @@ import org.thingsboard.server.common.data.id.EntityId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.plugin.ComponentLifecycleEvent;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
@@ -36,10 +35,10 @@ import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-class TbEntityLocalSubsInfoTest {
+public class TbEntityLocalSubsInfoTest {
 
     @Test
-    void addTest() {
+    public void addTest() {
         Set<TbAttributeSubscription> expectedSubs = new HashSet<>();
         TbEntityLocalSubsInfo subsInfo = createSubsInfo();
         TenantId tenantId = subsInfo.getTenantId();
@@ -74,7 +73,7 @@ class TbEntityLocalSubsInfoTest {
     }
 
     @Test
-    void removeTest() {
+    public void removeTest() {
         Set<TbAttributeSubscription> expectedSubs = new HashSet<>();
         TbEntityLocalSubsInfo subsInfo = createSubsInfo();
         TenantId tenantId = subsInfo.getTenantId();
@@ -117,8 +116,7 @@ class TbEntityLocalSubsInfoTest {
     }
 
     @Test
-    void removeAllTest() {
-        List<TbAttributeSubscription> subs = new ArrayList<>();
+    public void removeAllTest() {
         TbEntityLocalSubsInfo subsInfo = createSubsInfo();
         TenantId tenantId = subsInfo.getTenantId();
         EntityId entityId = subsInfo.getEntityId();
@@ -136,17 +134,28 @@ class TbEntityLocalSubsInfoTest {
                 .keyStates(Map.of("key3", 3L, "key4", 4L))
                 .build();
 
-        subs.add(attrSubscription1);
-        subs.add(attrSubscription2);
+        TbAttributeSubscription attrSubscription3 = TbAttributeSubscription.builder()
+                .sessionId("session3")
+                .tenantId(tenantId)
+                .entityId(entityId)
+                .keyStates(Map.of("key5", 5L, "key6", 6L))
+                .build();
 
         subsInfo.add(attrSubscription1);
         subsInfo.add(attrSubscription2);
+        subsInfo.add(attrSubscription3);
+
+        assertFalse(subsInfo.isEmpty());
+
+        TbEntitySubEvent updatedEvent = subsInfo.removeAll(List.of(attrSubscription1, attrSubscription2));
+        assertNotNull(updatedEvent);
+        checkEvent(updatedEvent, Set.of(attrSubscription3), ComponentLifecycleEvent.UPDATED);
 
         assertFalse(subsInfo.isEmpty());
 
-        TbEntitySubEvent deletedEvent = subsInfo.removeAll(subs);
+        TbEntitySubEvent deletedEvent = subsInfo.removeAll(List.of(attrSubscription3));
         assertNotNull(deletedEvent);
-        checkEvent(deletedEvent, subs, ComponentLifecycleEvent.DELETED);
+        checkEvent(deletedEvent, null, ComponentLifecycleEvent.DELETED);
 
         assertTrue(subsInfo.isEmpty());
     }
@@ -155,7 +164,7 @@ class TbEntityLocalSubsInfoTest {
         return new TbEntityLocalSubsInfo(new TenantId(UUID.randomUUID()), new DeviceId(UUID.randomUUID()));
     }
 
-    private void checkEvent(TbEntitySubEvent event, Collection<TbAttributeSubscription> expectedSubs, ComponentLifecycleEvent expectedType) {
+    private void checkEvent(TbEntitySubEvent event, Set<TbAttributeSubscription> expectedSubs, ComponentLifecycleEvent expectedType) {
         assertEquals(expectedType, event.getType());
         TbSubscriptionsInfo info = event.getInfo();
         if (event.getType() == ComponentLifecycleEvent.DELETED) {
@@ -171,7 +180,7 @@ class TbEntityLocalSubsInfoTest {
         assertEquals(getAttrKeys(expectedSubs), info.attrKeys);
     }
 
-    private Set<String> getAttrKeys(Collection<TbAttributeSubscription> attributeSubscriptions) {
+    private Set<String> getAttrKeys(Set<TbAttributeSubscription> attributeSubscriptions) {
         return attributeSubscriptions.stream().map(s -> s.getKeyStates().keySet()).flatMap(Collection::stream).collect(Collectors.toSet());
     }
 }