diff --git a/util/src/test/java/io/grpc/util/MultiChildLoadBalancerTest.java b/util/src/test/java/io/grpc/util/MultiChildLoadBalancerTest.java index fd47261539..d90c5eab92 100644 --- a/util/src/test/java/io/grpc/util/MultiChildLoadBalancerTest.java +++ b/util/src/test/java/io/grpc/util/MultiChildLoadBalancerTest.java @@ -337,14 +337,27 @@ public class MultiChildLoadBalancerTest { } private class TestLb extends MultiChildLoadBalancer { - protected TestLb(Helper mockHelper) { super(mockHelper); } @Override - protected SubchannelPicker getSubchannelPicker(Map childPickers) { - return new TestSubchannelPicker(childPickers); + protected void updateOverallBalancingState() { + ConnectivityState overallState = null; + final Map childPickers = new HashMap<>(); + for (ChildLbState childLbState : getChildLbStates()) { + if (childLbState.isDeactivated()) { + continue; + } + childPickers.put(childLbState.getKey(), childLbState.getCurrentPicker()); + overallState = aggregateState(overallState, childLbState.getCurrentState()); + } + + if (overallState != null) { + getHelper().updateBalancingState(overallState, new TestSubchannelPicker(childPickers)); + currentConnectivityState = overallState; + } + } private class TestSubchannelPicker extends SubchannelPicker {