xds: improve PriorityLoadBalancerTest

grpc · Mar 2, 2022 · 421be66 · 421be66
1 parent 1b1b284
commit 421be66
Showing 1 changed file with 14 additions and 44 deletions.
diff --git a/xds/src/test/java/io/grpc/xds/PriorityLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/PriorityLoadBalancerTest.java
@@ -425,16 +425,10 @@ public void typicalPriorityFailOverFlowWithIdleUpdate() {
     Helper helper0 = Iterables.getOnlyElement(fooHelpers);
 
     // p0 gets IDLE.
-    final Subchannel subchannel0 = mock(Subchannel.class);
     helper0.updateBalancingState(
         IDLE,
-        new SubchannelPicker() {
-          @Override
-          public PickResult pickSubchannel(PickSubchannelArgs args) {
-            return PickResult.withSubchannel(subchannel0);
-          }
-        });
-    assertCurrentPickerPicksIdleSubchannel(subchannel0);
+        BUFFER_PICKER);
+    assertCurrentPickerIsBufferPicker();
 
     // p0 fails over to p1 immediately.
     helper0.updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(Status.ABORTED));
@@ -452,32 +446,20 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
     Helper helper2 = Iterables.getLast(fooHelpers);
 
     // p2 gets IDLE
-    final Subchannel subchannel1 = mock(Subchannel.class);
     helper2.updateBalancingState(
         IDLE,
-        new SubchannelPicker() {
-          @Override
-          public PickResult pickSubchannel(PickSubchannelArgs args) {
-            return PickResult.withSubchannel(subchannel1);
-          }
-        });
-    assertCurrentPickerPicksIdleSubchannel(subchannel1);
+        BUFFER_PICKER);
+    assertCurrentPickerIsBufferPicker();
 
     // p0 gets back to IDLE
-    final Subchannel subchannel2 = mock(Subchannel.class);
     helper0.updateBalancingState(
         IDLE,
-        new SubchannelPicker() {
-          @Override
-          public PickResult pickSubchannel(PickSubchannelArgs args) {
-            return PickResult.withSubchannel(subchannel2);
-          }
-        });
-    assertCurrentPickerPicksIdleSubchannel(subchannel2);
+        BUFFER_PICKER);
+    assertCurrentPickerIsBufferPicker();
 
     // p2 fails but does not affect overall picker
     helper2.updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(Status.UNAVAILABLE));
-    assertCurrentPickerPicksIdleSubchannel(subchannel2);
+    assertCurrentPickerIsBufferPicker();
 
     // p0 fails over to p3 immediately since p1 already timeout and p2 already in TRANSIENT_FAILURE.
     helper0.updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(Status.UNAVAILABLE));
@@ -497,32 +479,20 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
     assertCurrentPickerReturnsError(Status.Code.DATA_LOSS, "foo");
 
     // p2 gets back to IDLE
-    final Subchannel subchannel3 = mock(Subchannel.class);
     helper2.updateBalancingState(
         IDLE,
-        new SubchannelPicker() {
-          @Override
-          public PickResult pickSubchannel(PickSubchannelArgs args) {
-            return PickResult.withSubchannel(subchannel3);
-          }
-        });
-    assertCurrentPickerPicksIdleSubchannel(subchannel3);
+        BUFFER_PICKER);
+    assertCurrentPickerIsBufferPicker();
 
     // p0 gets back to IDLE
-    final Subchannel subchannel4 = mock(Subchannel.class);
     helper0.updateBalancingState(
         IDLE,
-        new SubchannelPicker() {
-          @Override
-          public PickResult pickSubchannel(PickSubchannelArgs args) {
-            return PickResult.withSubchannel(subchannel4);
-          }
-        });
-    assertCurrentPickerPicksIdleSubchannel(subchannel4);
+        BUFFER_PICKER);
+    assertCurrentPickerIsBufferPicker();
 
     // p0 fails over to p2 and picker is updated to p2's existing picker.
     helper0.updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(Status.UNAVAILABLE));
-    assertCurrentPickerPicksIdleSubchannel(subchannel3);
+    assertCurrentPickerIsBufferPicker();
 
     // Deactivate child balancer get deleted.
     fakeClock.forwardTime(15, TimeUnit.MINUTES);
@@ -607,9 +577,9 @@ private void assertCurrentPickerPicksSubchannel(Subchannel expectedSubchannelToP
     assertThat(pickResult.getSubchannel()).isEqualTo(expectedSubchannelToPick);
   }
 
-  private void assertCurrentPickerPicksIdleSubchannel(Subchannel expectedSubchannelToPick) {
+  private void assertCurrentPickerIsBufferPicker() {
     assertLatestConnectivityState(IDLE);
     PickResult pickResult = pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class));
-    assertThat(pickResult.getSubchannel()).isEqualTo(expectedSubchannelToPick);
+    assertThat(pickResult).isEqualTo(PickResult.withNoResult());
   }
 }