Skip to content

Commit 420a6df

Browse files
voidzcyvoidzcy
committed
added unit tests
1 parent cced047 commit 420a6df

File tree

2 files changed

+193
-8
lines changed

2 files changed

+193
-8
lines changed

‎xds/src/main/java/io/grpc/xds/OrcaUtil.java‎

Lines changed: 29 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package io.grpc.xds;
1818

19+
import com.google.common.annotations.VisibleForTesting;
1920
import io.envoyproxy.udpa.data.orca.v1.OrcaLoadReport;
2021
import io.grpc.Attributes;
2122
import io.grpc.CallOptions;
@@ -26,6 +27,7 @@
2627
import io.grpc.protobuf.ProtoUtils;
2728
import io.grpc.util.ForwardingClientStreamTracer;
2829
import java.util.ArrayList;
30+
import java.util.Collections;
2931
import java.util.List;
3032
import java.util.concurrent.TimeUnit;
3133

@@ -53,7 +55,7 @@ public ClientStreamTracer newClientStreamTracer(StreamInfo info, Metadata header
5355
*/
5456
public static ClientStreamTracer.Factory newOrcaClientStreamTracerFactory(
5557
OrcaReportListener listener) {
56-
return newOrcaClientStreamTracerFacotry(NOOP_CLIENT_STREAM_TRACER_FACTORY, listener);
58+
return newOrcaClientStreamTracerFactory(NOOP_CLIENT_STREAM_TRACER_FACTORY, listener);
5759
}
5860

5961
/**
@@ -64,7 +66,7 @@ public static ClientStreamTracer.Factory newOrcaClientStreamTracerFactory(
6466
* @param listener contains the callback to be invoked when an per-request ORCA report is
6567
* received.
6668
*/
67-
public static ClientStreamTracer.Factory newOrcaClientStreamTracerFacotry(
69+
public static ClientStreamTracer.Factory newOrcaClientStreamTracerFactory(
6870
ClientStreamTracer.Factory delegate, OrcaReportListener listener) {
6971
return new OrcaClientStreamTracerFactory(delegate, listener);
7072
}
@@ -104,13 +106,16 @@ public interface OrcaReportListener {
104106
void onLoadReport(OrcaLoadReport report);
105107
}
106108

107-
private static final class OrcaClientStreamTracerFactory extends ClientStreamTracer.Factory {
109+
@VisibleForTesting
110+
static final class OrcaClientStreamTracerFactory extends ClientStreamTracer.Factory {
108111

109-
private static final CallOptions.Key<OrcaReportBroker> ORCA_REPORT_BROKER_KEY =
112+
@VisibleForTesting
113+
static final CallOptions.Key<OrcaReportBroker> ORCA_REPORT_BROKER_KEY =
110114
CallOptions.Key.create("internal-orca-report-broker");
111-
private static final Metadata.Key<OrcaLoadReport> ORCA_ENDPOINT_LOAD_METRICS_KEY =
115+
@VisibleForTesting
116+
static final Metadata.Key<OrcaLoadReport> ORCA_ENDPOINT_LOAD_METRICS_KEY =
112117
Metadata.Key.of(
113-
"X-Endpoint-Load-Metrics-Bin",
118+
"x-endpoint-load-metrics-bin",
114119
ProtoUtils.metadataMarshaller(OrcaLoadReport.getDefaultInstance()));
115120

116121
private final ClientStreamTracer.Factory delegate;
@@ -148,6 +153,7 @@ public CallOptions getCallOptions() {
148153
if (augmented) {
149154
final ClientStreamTracer currTracer = tracer;
150155
final OrcaReportBroker currBroker = broker;
156+
// The actual tracer that performs ORCA report deserialization.
151157
tracer = new ForwardingClientStreamTracer() {
152158
@Override
153159
protected ClientStreamTracer delegate() {
@@ -172,14 +178,29 @@ public void inboundTrailers(Metadata trailers) {
172178
* An {@link OrcaReportBroker} instance holds registered {@link OrcaReportListener}s and invoke
173179
* all of them when an {@link OrcaLoadReport} is received.
174180
*/
175-
private static final class OrcaReportBroker {
181+
@VisibleForTesting
182+
static final class OrcaReportBroker {
183+
184+
private final List<OrcaReportListener> listeners;
176185

177-
private final List<OrcaReportListener> listeners = new ArrayList<>();
186+
OrcaReportBroker() {
187+
this(new ArrayList<OrcaReportListener>());
188+
}
189+
190+
@VisibleForTesting
191+
OrcaReportBroker(List<OrcaReportListener> listeners) {
192+
this.listeners = listeners;
193+
}
178194

179195
void addListener(OrcaReportListener listener) {
180196
listeners.add(listener);
181197
}
182198

199+
@VisibleForTesting
200+
List<OrcaReportListener> getListeners() {
201+
return Collections.unmodifiableList(listeners);
202+
}
203+
183204
void onReport(OrcaLoadReport report) {
184205
for (OrcaReportListener listener : listeners) {
185206
listener.onLoadReport(report);

‎xds/src/test/java/io/grpc/xds/OrcaUtilTest.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*
2+
* Copyright 2019 The gRPC Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.grpc.xds;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
import static org.junit.Assert.assertEquals;
21+
import static org.junit.Assert.assertNotEquals;
22+
import static org.junit.Assert.assertSame;
23+
import static org.mockito.ArgumentMatchers.any;
24+
import static org.mockito.Mockito.doNothing;
25+
import static org.mockito.Mockito.mock;
26+
import static org.mockito.Mockito.never;
27+
import static org.mockito.Mockito.spy;
28+
import static org.mockito.Mockito.verify;
29+
import static org.mockito.Mockito.when;
30+
31+
import io.envoyproxy.udpa.data.orca.v1.OrcaLoadReport;
32+
import io.grpc.Attributes;
33+
import io.grpc.CallOptions;
34+
import io.grpc.ClientStreamTracer;
35+
import io.grpc.ClientStreamTracer.StreamInfo;
36+
import io.grpc.Metadata;
37+
import io.grpc.xds.OrcaUtil.OrcaReportBroker;
38+
import io.grpc.xds.OrcaUtil.OrcaReportListener;
39+
import org.junit.Test;
40+
import org.junit.runner.RunWith;
41+
import org.junit.runners.JUnit4;
42+
import org.mockito.ArgumentCaptor;
43+
44+
/** Unit tests for {@link OrcaUtil}. */
45+
@RunWith(JUnit4.class)
46+
public class OrcaUtilTest {
47+
48+
private static final ClientStreamTracer.StreamInfo STREAM_INFO =
49+
new ClientStreamTracer.StreamInfo() {
50+
@Override
51+
public Attributes getTransportAttrs() {
52+
return Attributes.EMPTY;
53+
}
54+
55+
@Override
56+
public CallOptions getCallOptions() {
57+
return CallOptions.DEFAULT;
58+
}
59+
};
60+
61+
/**
62+
* Test a single load balance policy's listener receive per-request ORCA reports upon call trailer
63+
* arrives.
64+
*/
65+
@Test
66+
public void singlePolicyPerRequestListener() {
67+
OrcaReportListener mockListener = mock(OrcaReportListener.class);
68+
// Use a mocked noop stream tracer factory.
69+
ClientStreamTracer.Factory fakeDelegateFactory = mock(ClientStreamTracer.Factory.class);
70+
ClientStreamTracer fakeTracer = mock(ClientStreamTracer.class);
71+
doNothing().when(fakeTracer).inboundTrailers(any(Metadata.class));
72+
when(fakeDelegateFactory.newClientStreamTracer(
73+
any(ClientStreamTracer.StreamInfo.class), any(Metadata.class)))
74+
.thenReturn(fakeTracer);
75+
76+
// The OrcaClientStreamTracerFactory will augment the StreamInfo passed to its
77+
// newClientStreamTracer method. The augmented StreamInfo's CallOptions will contain
78+
// a OrcaReportBroker, in which has the registered listener.
79+
ClientStreamTracer.Factory factory =
80+
OrcaUtil.newOrcaClientStreamTracerFactory(fakeDelegateFactory, mockListener);
81+
ClientStreamTracer tracer = factory.newClientStreamTracer(STREAM_INFO, new Metadata());
82+
ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = ArgumentCaptor.forClass(null);
83+
verify(fakeDelegateFactory)
84+
.newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class));
85+
ClientStreamTracer.StreamInfo capturedInfo = streamInfoCaptor.getValue();
86+
assertNotEquals(STREAM_INFO, capturedInfo);
87+
OrcaReportBroker broker =
88+
capturedInfo
89+
.getCallOptions()
90+
.getOption(OrcaUtil.OrcaClientStreamTracerFactory.ORCA_REPORT_BROKER_KEY);
91+
assertThat(broker).isNotNull();
92+
assertThat(broker.getListeners()).containsExactly(mockListener);
93+
94+
// When the trailer does not contain ORCA report, listener callback will not be invoked.
95+
Metadata trailer = new Metadata();
96+
tracer.inboundTrailers(trailer);
97+
verify(mockListener, never()).onLoadReport(any(OrcaLoadReport.class));
98+
99+
// When the trailer contains an ORCA report, listener callback will be invoked.
100+
trailer.put(
101+
OrcaUtil.OrcaClientStreamTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY,
102+
OrcaLoadReport.getDefaultInstance());
103+
tracer.inboundTrailers(trailer);
104+
ArgumentCaptor<OrcaLoadReport> reportCaptor = ArgumentCaptor.forClass(null);
105+
verify(mockListener).onLoadReport(reportCaptor.capture());
106+
assertEquals(OrcaLoadReport.getDefaultInstance(), reportCaptor.getValue());
107+
}
108+
109+
/**
110+
* Test parent-child load balance policies' listener both receive per-request ORCA reports upon
111+
* call trailer arrives and ORCA report deserialization happens only once.
112+
*/
113+
@Test
114+
public void twoLevelPolicyPerRequestListeners() {
115+
// Wrap the child factory for the sake of spy.
116+
OrcaReportListener childListener = mock(OrcaReportListener.class);
117+
final ClientStreamTracer.Factory childFactory =
118+
OrcaUtil.newOrcaClientStreamTracerFactory(childListener);
119+
ClientStreamTracer.Factory forwardingChildFactory =
120+
new ClientStreamTracer.Factory() {
121+
@Override
122+
public ClientStreamTracer newClientStreamTracer(StreamInfo info, Metadata headers) {
123+
return childFactory.newClientStreamTracer(info, headers);
124+
}
125+
};
126+
ClientStreamTracer.Factory spyChildFactory = spy(forwardingChildFactory);
127+
128+
OrcaReportListener parentListener = mock(OrcaReportListener.class);
129+
ClientStreamTracer.Factory parentFactory =
130+
OrcaUtil.newOrcaClientStreamTracerFactory(spyChildFactory, parentListener);
131+
ClientStreamTracer parentTracer =
132+
parentFactory.newClientStreamTracer(STREAM_INFO, new Metadata());
133+
ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = ArgumentCaptor.forClass(null);
134+
verify(spyChildFactory).newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class));
135+
ClientStreamTracer.StreamInfo childStreamInfo = streamInfoCaptor.getValue();
136+
137+
assertNotEquals(childStreamInfo, STREAM_INFO);
138+
OrcaReportBroker broker =
139+
childStreamInfo
140+
.getCallOptions()
141+
.getOption(OrcaUtil.OrcaClientStreamTracerFactory.ORCA_REPORT_BROKER_KEY);
142+
assertThat(broker.getListeners()).containsExactly(parentListener, childListener);
143+
144+
// When the trailer does not contain ORCA report, no listener callback will be invoked.
145+
Metadata trailer = new Metadata();
146+
parentTracer.inboundTrailers(trailer);
147+
verify(parentListener, never()).onLoadReport(any(OrcaLoadReport.class));
148+
verify(childListener, never()).onLoadReport(any(OrcaLoadReport.class));
149+
150+
// When the trailer contains an ORCA report, callbacks for both listeners will be invoked.
151+
// Both listener will receive the same ORCA report instance, which means deserialization
152+
// happens only once.
153+
trailer.put(
154+
OrcaUtil.OrcaClientStreamTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY,
155+
OrcaLoadReport.getDefaultInstance());
156+
parentTracer.inboundTrailers(trailer);
157+
ArgumentCaptor<OrcaLoadReport> parentReportCap = ArgumentCaptor.forClass(null);
158+
ArgumentCaptor<OrcaLoadReport> childReportCap = ArgumentCaptor.forClass(null);
159+
verify(parentListener).onLoadReport(parentReportCap.capture());
160+
verify(childListener).onLoadReport(childReportCap.capture());
161+
assertEquals(OrcaLoadReport.getDefaultInstance(), parentReportCap.getValue());
162+
assertSame(parentReportCap.getValue(), childReportCap.getValue());
163+
}
164+
}

0 commit comments

Comments
 (0)