diff --git a/zuul-core/src/main/java/com/netflix/zuul/netty/server/push/PushConnectionRegistry.java b/zuul-core/src/main/java/com/netflix/zuul/netty/server/push/PushConnectionRegistry.java index 1975294419..da79a75d63 100644 --- a/zuul-core/src/main/java/com/netflix/zuul/netty/server/push/PushConnectionRegistry.java +++ b/zuul-core/src/main/java/com/netflix/zuul/netty/server/push/PushConnectionRegistry.java @@ -15,13 +15,15 @@ */ package com.netflix.zuul.netty.server.push; -import javax.annotation.Nullable; -import javax.inject.Inject; -import javax.inject.Singleton; import java.security.SecureRandom; +import java.util.ArrayList; import java.util.Base64; +import java.util.List; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import javax.annotation.Nullable; +import javax.inject.Inject; +import javax.inject.Singleton; /** * Maintains client identity to web socket or SSE channel mapping. @@ -45,6 +47,10 @@ public PushConnection get(final String clientId) { return clientPushConnectionMap.get(clientId); } + public List getAll() { + return new ArrayList<>(clientPushConnectionMap.values()); + } + public String mintNewSecureToken() { byte[] tokenBuffer = new byte[15]; secureTokenGenerator.nextBytes(tokenBuffer); diff --git a/zuul-core/src/test/java/com/netflix/zuul/netty/server/push/PushConnectionRegistryTest.java b/zuul-core/src/test/java/com/netflix/zuul/netty/server/push/PushConnectionRegistryTest.java new file mode 100644 index 0000000000..a3b7d847c3 --- /dev/null +++ b/zuul-core/src/test/java/com/netflix/zuul/netty/server/push/PushConnectionRegistryTest.java @@ -0,0 +1,74 @@ +package com.netflix.zuul.netty.server.push; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class PushConnectionRegistryTest { + private PushConnectionRegistry pushConnectionRegistry; + + private PushConnection pushConnection; + + @BeforeEach + void setUp() { + pushConnectionRegistry = new PushConnectionRegistry(); + pushConnection = mock(PushConnection.class); + } + + @Test + void testPutAndGet() { + assertNull(pushConnectionRegistry.get("clientId1")); + + pushConnectionRegistry.put("clientId1", pushConnection); + + assertEquals(pushConnection, pushConnectionRegistry.get("clientId1")); + } + + @Test + void testGetAll() { + pushConnectionRegistry.put("clientId1", pushConnection); + pushConnectionRegistry.put("clientId2", pushConnection); + + List connections = pushConnectionRegistry.getAll(); + + assertEquals(2, connections.size()); + } + + @Test + void testMintNewSecureToken() { + String token = pushConnectionRegistry.mintNewSecureToken(); + + assertNotNull(token); + assertEquals(20, token.length()); // 15 bytes become 20 characters when Base64-encoded + } + + @Test + void testPutAssignsTokenToConnection() { + pushConnectionRegistry.put("clientId1", pushConnection); + + verify(pushConnection).setSecureToken(anyString()); + } + + @Test + void testRemove() { + pushConnectionRegistry.put("clientId1", pushConnection); + + assertEquals(pushConnection, pushConnectionRegistry.remove("clientId1")); + assertNull(pushConnectionRegistry.get("clientId1")); + } + + @Test + void testSize() { + assertEquals(0, pushConnectionRegistry.size()); + + pushConnectionRegistry.put("clientId1", pushConnection); + + assertEquals(1, pushConnectionRegistry.size()); + } +}