util: duplicate RoundRobinLoadBalancerTest.pickAfterStateChange for new pick-first

PickFirstLeafLoadBalancer does not trigger a second refreshNameResolution
when transitioning from TRANSIENT_FAILURE to IDLE during backoff.
Duplicate pickAfterStateChange to verify both old and new PF behaviors.

TAG=agy
CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc

Author: kannanjgithub

util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java

@@ -230,6 +230,7 @@ public void pickAfterResolvedUpdatedHosts() throws Exception {
 
   @Test
   public void pickAfterStateChange() throws Exception {
+    PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false);
     InOrder inOrder = inOrder(mockHelper);
     Status addressesAcceptanceStatus =
         acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY);
@@ -263,6 +264,42 @@ public void pickAfterStateChange() throws Exception {
     AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper);
   }
 
+  @Test
+  public void pickAfterStateChange_newPickFirst() throws Exception {
+    PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true);
+    InOrder inOrder = inOrder(mockHelper);
+    Status addressesAcceptanceStatus =
+        acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY);
+    assertThat(addressesAcceptanceStatus.isOk()).isTrue();
+    inOrder.verify(mockHelper).createSubchannel(any(CreateSubchannelArgs.class));
+
+    // TODO figure out if this method testing the right things
+
+    assertThat(subchannels).hasSize(1);
+    Subchannel subchannel = subchannels.values().iterator().next();
+
+    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER));
+
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class);
+
+    Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯");
+    deliverSubchannelState(subchannel,
+        ConnectivityStateInfo.forTransientFailure(error));
+    AbstractTestHelper.refreshInvokedAndUpdateBS(
+        inOrder, TRANSIENT_FAILURE, mockHelper, pickerCaptor);
+    assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error);
+
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
+    inOrder.verify(mockHelper, never()).refreshNameResolution();
+    inOrder.verify(mockHelper, never())
+        .updateBalancingState(eq(TRANSIENT_FAILURE), any(SubchannelPicker.class));
+
+    verify(subchannel, atLeastOnce()).requestConnection();
+    AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper);
+  }
+
   @Test
   public void ignoreShutdownSubchannelStateChange() {
     InOrder inOrder = inOrder(mockHelper);