HBASE-26648 Improve fidelity of RegionLocator spans (addendum)

Addendum extends the test case to cover both `RpcConnectionRegistry` and `ZKConnectionRegistry`.

hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncMetaRegionLocator.java

@@ -34,7 +34,9 @@ import java.util.concurrent.TimeUnit;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.ConnectionRule;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
+import org.apache.hadoop.hbase.HBaseConfiguration;
 import org.apache.hadoop.hbase.HBaseTestingUtil;
+import org.apache.hadoop.hbase.HConstants;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.MatcherPredicate;
 import org.apache.hadoop.hbase.MiniClusterRule;
@@ -54,39 +56,61 @@ import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.experimental.runners.Enclosed;
 import org.junit.rules.ExternalResource;
 import org.junit.rules.RuleChain;
+import org.junit.rules.TestName;
 import org.junit.rules.TestRule;
+import org.junit.runner.RunWith;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-@Category({ MediumTests.class, ClientTests.class })
+@RunWith(Enclosed.class)
 public class TestAsyncMetaRegionLocator {
   private static final Logger logger = LoggerFactory.getLogger(TestAsyncMetaRegionLocator.class);
 
-  @ClassRule
-  public static final HBaseClassTestRule CLASS_RULE =
-    HBaseClassTestRule.forClass(TestAsyncMetaRegionLocator.class);
-
-  private static final OpenTelemetryClassRule otelClassRule = OpenTelemetryClassRule.create();
-  private static final MiniClusterRule miniClusterRule = MiniClusterRule.newBuilder()
-    .setMiniClusterOption(StartTestingClusterOption.builder().numWorkers(3).build()).build();
-  private static final ConnectionRule connectionRule =
-    ConnectionRule.createAsyncConnectionRule(miniClusterRule::createAsyncConnection);
-
   private static final class Setup extends ExternalResource {
+
+    private final MiniClusterRule miniClusterRule;
+    private final ConnectionRule connectionRule;
+
+    private boolean initialized = false;
+    private HBaseTestingUtil testUtil;
+    private AsyncMetaRegionLocator locator;
     private ConnectionRegistry registry;
 
+    public Setup(final ConnectionRule connectionRule, final MiniClusterRule miniClusterRule) {
+      this.connectionRule = connectionRule;
+      this.miniClusterRule = miniClusterRule;
+    }
+
+    public HBaseTestingUtil getTestingUtility() {
+      assertInitialized();
+      return testUtil;
+    }
+
+    public AsyncMetaRegionLocator getLocator() {
+      assertInitialized();
+      return locator;
+    }
+
+    private void assertInitialized() {
+      if (!initialized) {
+        throw new IllegalStateException("before method has not been called.");
+      }
+    }
+
     @Override
     protected void before() throws Throwable {
       final AsyncAdmin admin = connectionRule.getAsyncConnection().getAdmin();
-      TEST_UTIL = miniClusterRule.getTestingUtility();
+      testUtil = miniClusterRule.getTestingUtility();
       HBaseTestingUtil.setReplicas(admin, TableName.META_TABLE_NAME, 3);
-      TEST_UTIL.waitUntilNoRegionsInTransition();
+      testUtil.waitUntilNoRegionsInTransition();
-      registry = ConnectionRegistryFactory.getRegistry(TEST_UTIL.getConfiguration());
+      registry = ConnectionRegistryFactory.getRegistry(testUtil.getConfiguration());
-      RegionReplicaTestHelper.waitUntilAllMetaReplicasAreReady(TEST_UTIL, registry);
+      RegionReplicaTestHelper.waitUntilAllMetaReplicasAreReady(testUtil, registry);
       admin.balancerSwitch(false).get();
-      LOCATOR = new AsyncMetaRegionLocator(registry);
+      locator = new AsyncMetaRegionLocator(registry);
+      initialized = true;
     }
 
     @Override
@@ -95,62 +119,127 @@ public class TestAsyncMetaRegionLocator {
     }
   }
 
-  @ClassRule
+  public static abstract class AbstractBase {
-  public static final TestRule classRule = RuleChain.outerRule(otelClassRule)
+    private final OpenTelemetryClassRule otelClassRule = OpenTelemetryClassRule.create();
-    .around(miniClusterRule).around(connectionRule).around(new Setup());
+    private final MiniClusterRule miniClusterRule;
+    private final Setup setup;
 
-  private static HBaseTestingUtil TEST_UTIL;
+    protected Matcher<SpanData> parentSpanMatcher;
-  private static AsyncMetaRegionLocator LOCATOR;
+    protected List<SpanData> spans;
+    protected Matcher<SpanData> registryGetMetaRegionLocationsMatcher;
+
+    @Rule
+    public final TestRule classRule;
 
     @Rule
     public final OpenTelemetryTestRule otelTestRule = new OpenTelemetryTestRule(otelClassRule);
 
+    @Rule
+    public TestName testName = new TestName();
+
+    public AbstractBase() {
+      miniClusterRule = MiniClusterRule.newBuilder()
+        .setMiniClusterOption(StartTestingClusterOption.builder().numWorkers(3).build())
+        .setConfiguration(() -> {
+          final Configuration conf = HBaseConfiguration.create();
+          conf.setClass(HConstants.CLIENT_CONNECTION_REGISTRY_IMPL_CONF_KEY,
+            getConnectionRegistryClass(), ConnectionRegistry.class);
+          return conf;
+        }).build();
+      final ConnectionRule connectionRule =
+        ConnectionRule.createAsyncConnectionRule(miniClusterRule::createAsyncConnection);
+      setup = new Setup(connectionRule, miniClusterRule);
+      classRule = RuleChain.outerRule(otelClassRule).around(miniClusterRule).around(connectionRule)
+        .around(setup);
+    }
+
+    protected abstract Class<? extends ConnectionRegistry> getConnectionRegistryClass();
+
     @Test
     public void test() throws Exception {
+      final AsyncMetaRegionLocator locator = setup.getLocator();
+      final HBaseTestingUtil testUtil = setup.getTestingUtility();
+
       TraceUtil.trace(() -> {
         try {
-        testLocator(miniClusterRule.getTestingUtility(), TableName.META_TABLE_NAME, new Locator() {
+          testLocator(miniClusterRule.getTestingUtility(), TableName.META_TABLE_NAME,
+            new Locator() {
               @Override
               public void updateCachedLocationOnError(HRegionLocation loc, Throwable error) {
-            LOCATOR.updateCachedLocationOnError(loc, error);
+                locator.updateCachedLocationOnError(loc, error);
               }
 
               @Override
               public RegionLocations getRegionLocations(TableName tableName, int replicaId,
                 boolean reload) throws Exception {
-            return LOCATOR.getRegionLocations(replicaId, reload).get();
+                return locator.getRegionLocations(replicaId, reload).get();
               }
             });
         } catch (Exception e) {
           throw new RuntimeException(e);
         }
-    }, "test");
+      }, testName.getMethodName());
 
-    final Configuration conf = TEST_UTIL.getConfiguration();
+      final Configuration conf = testUtil.getConfiguration();
-    final Matcher<SpanData> parentSpanMatcher = allOf(hasName("test"), hasEnded());
+      parentSpanMatcher = allOf(hasName(testName.getMethodName()), hasEnded());
       Waiter.waitFor(conf, TimeUnit.SECONDS.toMillis(5),
         new MatcherPredicate<>(otelClassRule::getSpans, hasItem(parentSpanMatcher)));
-    final List<SpanData> spans = otelClassRule.getSpans();
+      spans = otelClassRule.getSpans();
       if (logger.isDebugEnabled()) {
         StringTraceRenderer renderer = new StringTraceRenderer(spans);
         renderer.render(logger::debug);
       }
-
       assertThat(spans, hasItem(parentSpanMatcher));
-    final SpanData parentSpan =
+      final SpanData parentSpan = spans.stream().filter(parentSpanMatcher::matches).findAny()
-      spans.stream().filter(parentSpanMatcher::matches).findAny().orElseThrow(AssertionError::new);
+        .orElseThrow(AssertionError::new);
 
-    final Matcher<SpanData> registryGetMetaRegionLocationsMatcher =
+      registryGetMetaRegionLocationsMatcher =
         allOf(hasName(endsWith("ConnectionRegistry.getMetaRegionLocations")),
           hasParentSpanId(parentSpan), hasKind(SpanKind.INTERNAL), hasEnded());
       assertThat(spans, hasItem(registryGetMetaRegionLocationsMatcher));
+    }
+  }
+
+  /**
+   * Test covers when client is configured with {@link ZKConnectionRegistry}.
+   */
+  @Category({ MediumTests.class, ClientTests.class })
+  public static class TestZKConnectionRegistry extends AbstractBase {
+    @ClassRule
+    public static final HBaseClassTestRule CLASS_RULE =
+      HBaseClassTestRule.forClass(TestZKConnectionRegistry.class);
+
+    @Override
+    protected Class<? extends ConnectionRegistry> getConnectionRegistryClass() {
+      return ZKConnectionRegistry.class;
+    }
+  }
+
+  /**
+   * Test covers when client is configured with {@link RpcConnectionRegistry}.
+   */
+  @Category({ MediumTests.class, ClientTests.class })
+  public static class TestRpcConnectionRegistry extends AbstractBase {
+    @ClassRule
+    public static final HBaseClassTestRule CLASS_RULE =
+      HBaseClassTestRule.forClass(TestRpcConnectionRegistry.class);
+
+    @Override
+    protected Class<? extends ConnectionRegistry> getConnectionRegistryClass() {
+      return RpcConnectionRegistry.class;
+    }
+
+    @Test
+    @Override
+    public void test() throws Exception {
+      super.test();
       final SpanData registry_getMetaRegionLocationsSpan =
         spans.stream().filter(registryGetMetaRegionLocationsMatcher::matches).findAny()
           .orElseThrow(AssertionError::new);
-
+      final Matcher<SpanData> clientGetMetaRegionLocationsMatcher = allOf(
-    final Matcher<SpanData> clientGetMetaRegionLocationsMatcher =
+        hasName(endsWith("ClientMetaService/GetMetaRegionLocations")),
-      allOf(hasName(endsWith("ClientMetaService/GetMetaRegionLocations")),
         hasParentSpanId(registry_getMetaRegionLocationsSpan), hasKind(SpanKind.CLIENT), hasEnded());
       assertThat(spans, hasItem(clientGetMetaRegionLocationsMatcher));
     }
+  }
 }