add unit test

linkedin · Sep 15, 2024 · ae592f6 · ae592f6
1 parent b374683
commit ae592f6
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 5 deletions.
diff --git a/d2/src/main/java/com/linkedin/d2/xds/XdsClient.java b/d2/src/main/java/com/linkedin/d2/xds/XdsClient.java
@@ -103,7 +103,7 @@ public static abstract class WildcardResourceWatcher
      * Defining a private constructor means only classes that are defined in this file can extend this class (see
      * {@link ResourceWatcher}).
      */
-    private WildcardResourceWatcher(ResourceType type)
+    WildcardResourceWatcher(ResourceType type)
     {
       _type = type;
     }

diff --git a/d2/src/main/java/com/linkedin/d2/xds/XdsClientImpl.java b/d2/src/main/java/com/linkedin/d2/xds/XdsClientImpl.java
@@ -175,7 +175,7 @@ public void watchAllXdsResources(WildcardResourceWatcher watcher)
     _executorService.execute(() ->
     {
       _log.info("Subscribing to wildcard for resource type: {}", watcher.getType());
-      WildcardResourceSubscriber subscriber = _wildcardSubscribers.get(watcher.getType());
+      WildcardResourceSubscriber subscriber = getWildcardResourceSubscriber(watcher.getType());
       if (subscriber == null)
       {
         subscriber = new WildcardResourceSubscriber(watcher.getType());
@@ -568,6 +568,12 @@ Map<String, ResourceSubscriber> getResourceSubscriberMap(ResourceType type)
     return _resourceSubscribers.get(type);
   }
 
+  @VisibleForTesting
+  WildcardResourceSubscriber getWildcardResourceSubscriber(ResourceType type)
+  {
+    return _wildcardSubscribers.get(type);
+  }
+
   static class ResourceSubscriber
   {
     private final ResourceType _type;
@@ -840,7 +846,7 @@ public void run()
       for (ResourceType type : ResourceType.values())
       {
         Set<String> resources = new HashSet<>(getResourceSubscriberMap(type).keySet());
-        if (resources.isEmpty() && !_wildcardSubscribers.containsKey(type))
+        if (resources.isEmpty() && getWildcardResourceSubscriber(type) == null)
         {
           continue;
         }
@@ -856,7 +862,7 @@ public void run()
         {
           rewrittenType = type;
         }
-        if (_wildcardSubscribers.containsKey(type))
+        if (getWildcardResourceSubscriber(type) != null)
         {
           resources.add("*");
         }

diff --git a/d2/src/test/java/com/linkedin/d2/xds/TestXdsClientImpl.java b/d2/src/test/java/com/linkedin/d2/xds/TestXdsClientImpl.java
@@ -16,6 +16,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import io.grpc.Status;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
@@ -443,6 +444,70 @@ public void testHandleD2URICollectionResponseWithRemoval()
     Assert.assertEquals(actualData.getURIMap(), D2_URI_MAP_UPDATE_WITH_DATA1.getURIMap());
   }
 
+  @Test
+  public void testWildCardResourceSubscription()
+  {
+    XdsClientImplFixture fixture = new XdsClientImplFixture();
+    XdsClient.WildcardResourceWatcher nodeWildCardWatcher = new XdsClient.WildcardNodeResourceWatcher()
+    {
+      @Override public void onError(Status error)
+      {
+
+      }
+      @Override public void onReconnect()
+      {
+
+      }
+
+      @Override
+      public void onChanged(String resourceName, XdsClient.NodeUpdate nodeUpdate)
+      {
+
+      }
+
+      @Override public void onRemoval(String resourceName)
+      {
+
+      }
+    };
+
+
+    XdsClient.WildcardResourceWatcher uriMapWildCardWatcher = new XdsClient.WildcardD2URIMapResourceWatcher()
+    {
+      @Override public void onError(Status error)
+      {
+
+      }
+
+      @Override public void onReconnect()
+      {
+
+      }
+
+      @Override
+      public void onChanged(String resourceName, D2URIMapUpdate d2URIMapUpdate)
+      {
+      }
+
+      @Override public void onRemoval(String resourceName)
+      {
+
+      }
+
+    };
+
+    fixture._xdsClientImpl.getWildcardResourceSubscriber(NODE).addWatcher(nodeWildCardWatcher);
+    fixture._xdsClientImpl.getWildcardResourceSubscriber(D2_URI_MAP).addWatcher(uriMapWildCardWatcher);
+
+    fixture._xdsClientImpl.handleResponse(DISCOVERY_RESPONSE_NODE_DATA1);
+    fixture.verifyAckSent(1);
+    nodeWildCardWatcher.onChanged(SERVICE_RESOURCE_NAME , eq(NODE_UPDATE1));
+
+    fixture._xdsClientImpl.handleResponse(DISCOVERY_RESPONSE_URI_MAP_DATA1);
+    fixture.verifyAckSent(2);
+    uriMapWildCardWatcher.onChanged(CLUSTER_RESOURCE_NAME, eq(D2_URI_MAP_UPDATE_WITH_DATA1));
+  }
+
   @Test
   public void testResourceSubscriberAddWatcher()
   {
@@ -467,7 +532,11 @@ private static class XdsClientImplFixture
     XdsClientJmx _xdsClientJmx;
     ResourceSubscriber _nodeSubscriber;
     ResourceSubscriber _clusterSubscriber;
+    XdsClientImpl.WildcardResourceSubscriber _nodeWildcardSubscriber;
+    XdsClientImpl.WildcardResourceSubscriber _uriMapWildcardSubscriber;
     Map<ResourceType, Map<String, ResourceSubscriber>> _subscribers = new HashMap<>();
+    Map<ResourceType, XdsClientImpl.WildcardResourceSubscriber> _wildcardSubscribers = new HashMap<>();
+
     @Mock
     XdsClient.ResourceWatcher _resourceWatcher;
     @Mock
@@ -483,21 +552,30 @@ private static class XdsClientImplFixture
       MockitoAnnotations.initMocks(this);
       _nodeSubscriber = spy(new ResourceSubscriber(NODE, SERVICE_RESOURCE_NAME, _xdsClientJmx));
       _clusterSubscriber = spy(new ResourceSubscriber(D2_URI_MAP, CLUSTER_RESOURCE_NAME, _xdsClientJmx));
+      _nodeWildcardSubscriber = new XdsClientImpl.WildcardResourceSubscriber(NODE);
+      _uriMapWildcardSubscriber = new XdsClientImpl.WildcardResourceSubscriber(D2_URI_MAP);
+
+
 
       doNothing().when(_resourceWatcher).onChanged(any());
       for (ResourceSubscriber subscriber : Lists.newArrayList(_nodeSubscriber, _clusterSubscriber))
       {
         subscriber.addWatcher(_resourceWatcher);
         _subscribers.put(subscriber.getType(), Collections.singletonMap(subscriber.getResource(), subscriber));
       }
-
+      _wildcardSubscribers.put(NODE, _nodeWildcardSubscriber);
+      _wildcardSubscribers.put(D2_URI_MAP, _uriMapWildcardSubscriber);
       doNothing().when(_serverMetricsProvider).trackLatency(anyLong());
       _xdsClientImpl = spy(new XdsClientImpl(null, null, null, 0,
           useGlobCollections, _serverMetricsProvider));
       doNothing().when(_xdsClientImpl).sendAckOrNack(any(), any(), any());
       when(_xdsClientImpl.getXdsClientJmx()).thenReturn(_xdsClientJmx);
       when(_xdsClientImpl.getResourceSubscriberMap(any()))
           .thenAnswer(a -> _subscribers.get((ResourceType) a.getArguments()[0]));
+      doNothing().when(_xdsClientImpl).watchAllXdsResources(any());
+      when(_xdsClientImpl.getWildcardResourceSubscriber(any()))
+          .thenAnswer(a -> _wildcardSubscribers.get((ResourceType) a.getArguments()[0]));
+
     }
 
     void verifyAckSent(int count)