Skip to content

Commit 1c1922c

Browse files
parveensaniaparveensania
committed
Using mock in TC instead of java reflection
1 parent cd41fe9 commit 1c1922c

3 files changed

Lines changed: 93 additions & 29 deletions

File tree

‎runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -891,6 +891,11 @@ static StreamingDataflowWorker forTesting(
891891
createGrpcwindmillStreamFactoryBuilder(options, 1)
892892
.setProcessHeartbeatResponses(
893893
new WorkHeartbeatResponseProcessor(computationStateCache::get));
894+
GrpcDispatcherClient grpcDispatcherClient = GrpcDispatcherClient.create(options, stubFactory);
895+
grpcDispatcherClient.consumeWindmillDispatcherEndpoints(
896+
ImmutableSet.<HostAndPort>builder()
897+
.add(HostAndPort.fromHost("StreamingDataflowWorkerTest"))
898+
.build());
894899

895900
return new StreamingDataflowWorker(
896901
windmillServer,
@@ -916,7 +921,7 @@ static StreamingDataflowWorker forTesting(
916921
: windmillStreamFactory.build(),
917922
executorSupplier.apply("RefreshWork"),
918923
stageInfo,
919-
GrpcDispatcherClient.create(options, stubFactory));
924+
grpcDispatcherClient);
920925
}
921926

922927
private static GrpcWindmillStreamFactory.Builder createGrpcwindmillStreamFactoryBuilder(

‎runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
4848
import org.apache.beam.runners.dataflow.worker.streaming.WorkHeartbeatResponseProcessor;
4949
import org.apache.beam.runners.dataflow.worker.streaming.WorkId;
50+
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillMetadataServiceV1Alpha1Grpc;
5051
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
5152
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitWorkResponse;
5253
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationCommitWorkRequest;
@@ -60,12 +61,15 @@
6061
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution;
6162
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution.State;
6263
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest;
64+
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
65+
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
6366
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub;
6467
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
6568
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
6669
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
6770
import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemReceiver;
6871
import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
72+
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
6973
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
7074
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
7175
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
@@ -92,6 +96,8 @@ public final class FakeWindmillServer extends WindmillServerStub {
9296
private final ConcurrentHashMap<Long, Consumer<Windmill.CommitStatus>> droppedStreamingCommits;
9397
private final List<Windmill.GetDataRequest> getDataRequests = new ArrayList<>();
9498
private final Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses;
99+
private StreamObserver<WorkerMetadataResponse> workerMetadataObserver = null;
100+
95101
private int commitsRequested = 0;
96102
private boolean dropStreamingCommits = false;
97103

@@ -553,6 +559,47 @@ public synchronized void setWindmillServiceEndpoints(Set<HostAndPort> endpoints)
553559
this.dispatcherEndpoints = ImmutableSet.copyOf(endpoints);
554560
}
555561

562+
public void injectWorkerMetadata(WorkerMetadataResponse response) {
563+
if (workerMetadataObserver != null) {
564+
workerMetadataObserver.onNext(response);
565+
}
566+
}
567+
568+
private void setWorkerMetadataObserver(
569+
StreamObserver<WorkerMetadataResponse> workerMetadataObserver) {
570+
this.workerMetadataObserver = workerMetadataObserver;
571+
}
572+
573+
public static class FakeWindmillMetadataService
574+
extends CloudWindmillMetadataServiceV1Alpha1Grpc
575+
.CloudWindmillMetadataServiceV1Alpha1ImplBase {
576+
private final FakeWindmillServer server;
577+
578+
public FakeWindmillMetadataService(FakeWindmillServer server) {
579+
this.server = server;
580+
}
581+
582+
@Override
583+
public StreamObserver<WorkerMetadataRequest> getWorkerMetadata(
584+
StreamObserver<WorkerMetadataResponse> responseObserver) {
585+
server.setWorkerMetadataObserver(responseObserver);
586+
return new StreamObserver<WorkerMetadataRequest>() {
587+
@Override
588+
public void onNext(WorkerMetadataRequest value) {}
589+
590+
@Override
591+
public void onError(Throwable t) {
592+
responseObserver.onError(t);
593+
}
594+
595+
@Override
596+
public void onCompleted() {
597+
responseObserver.onCompleted();
598+
}
599+
};
600+
}
601+
}
602+
556603
public static class ResponseQueue<T, U> {
557604
private final Queue<Function<T, U>> responses = new ConcurrentLinkedQueue<>();
558605
Duration sleep = Duration.ZERO;

‎runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java‎

Lines changed: 40 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@
114114
import org.apache.beam.runners.dataflow.worker.testing.TestCountingSource;
115115
import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
116116
import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames;
117-
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillMetadataServiceV1Alpha1Grpc;
117+
import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
118118
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
119119
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus;
120120
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest;
@@ -136,10 +136,9 @@
136136
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer.Type;
137137
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WatermarkHold;
138138
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest;
139-
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
140139
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
140+
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
141141
import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient;
142-
import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcDispatcherClient;
143142
import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels;
144143
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactory;
145144
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactoryFactory;
@@ -190,13 +189,15 @@
190189
import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode;
191190
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
192191
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat;
193-
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Channel;
194-
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ClientCall;
192+
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server;
193+
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.inprocess.InProcessServerBuilder;
194+
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule;
195195
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats;
196196
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
197197
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
198198
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
199199
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
200+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
200201
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedLong;
201202
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
202203
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles;
@@ -294,6 +295,7 @@ public Long get() {
294295
@Rule public transient Timeout globalTimeout = Timeout.seconds(600);
295296
@Rule public BlockingFn blockingFn = new BlockingFn();
296297
@Rule public TestRule restoreMDC = new RestoreDataflowLoggingMDC();
298+
@Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
297299
@Rule public ErrorCollector errorCollector = new ErrorCollector();
298300
WorkUnitClient mockWorkUnitClient = mock(WorkUnitClient.class);
299301
StreamingGlobalConfigHandleImpl mockGlobalConfigHandle =
@@ -4074,6 +4076,19 @@ public void testSwitchStreamingWorkerHarness() throws Exception {
40744076
return;
40754077
}
40764078

4079+
String serverName = "StreamingDataflowWorkerTestChannel";
4080+
Server fakeServer =
4081+
grpcCleanup
4082+
.register(
4083+
InProcessServerBuilder.forName(serverName)
4084+
.directExecutor()
4085+
.addService(new FakeWindmillServer.FakeWindmillMetadataService(server))
4086+
.addService(
4087+
new CloudWindmillServiceV1Alpha1Grpc
4088+
.CloudWindmillServiceV1Alpha1ImplBase() {})
4089+
.build())
4090+
.start();
4091+
40774092
List<ParallelInstruction> instructions =
40784093
Arrays.asList(
40794094
makeSourceInstruction(StringUtf8Coder.of()),
@@ -4082,6 +4097,7 @@ public void testSwitchStreamingWorkerHarness() throws Exception {
40824097
// Start with CloudPath.
40834098
DataflowWorkerHarnessOptions options =
40844099
createTestingPipelineOptions("--isWindmillServiceDirectPathEnabled=false");
4100+
options.setWindmillServiceEndpoint(serverName);
40854101

40864102
StreamingDataflowWorker worker =
40874103
makeWorker(
@@ -4091,31 +4107,13 @@ public void testSwitchStreamingWorkerHarness() throws Exception {
40914107
.publishCounters()
40924108
.build());
40934109

4094-
GrpcDispatcherClient mockDispatcherClient = mock(GrpcDispatcherClient.class);
4095-
4096-
// FanOutStreamingEngineWorkerHarness creates
4097-
// CloudWindmillMetadataServiceV1Alpha1Stub and expects the stream to
4098-
// successfully start. Mocking it here.
4099-
Channel mockChannel = mock(Channel.class);
4100-
ClientCall<WorkerMetadataRequest, WorkerMetadataResponse> mockClientCall =
4101-
mock(ClientCall.class);
4102-
when(mockChannel.newCall(
4103-
eq(CloudWindmillMetadataServiceV1Alpha1Grpc.getGetWorkerMetadataMethod()), any()))
4104-
.thenReturn(mockClientCall);
4105-
when(mockDispatcherClient.getWindmillMetadataServiceStubBlocking())
4106-
.thenReturn(CloudWindmillMetadataServiceV1Alpha1Grpc.newStub(mockChannel));
4107-
java.lang.reflect.Field dispatcherClientField =
4108-
StreamingDataflowWorker.class.getDeclaredField("dispatcherClient");
4109-
dispatcherClientField.setAccessible(true);
4110-
dispatcherClientField.set(worker, mockDispatcherClient);
4111-
4112-
// Capture the config observer.
41134110
ArgumentCaptor<Consumer<StreamingGlobalConfig>> observerCaptor =
41144111
ArgumentCaptor.forClass(Consumer.class);
41154112

41164113
worker.start();
41174114

41184115
verify(mockGlobalConfigHandle, atLeastOnce()).registerConfigObserver(observerCaptor.capture());
4116+
41194117
List<Consumer<StreamingGlobalConfig>> observers = observerCaptor.getAllValues();
41204118

41214119
assertTrue(
@@ -4130,6 +4128,20 @@ public void testSwitchStreamingWorkerHarness() throws Exception {
41304128

41314129
ExecutorService harnessSwitchExecutor = worker.getHarnessSwitchExecutor();
41324130

4131+
// Prepare WorkerMetadataResponse
4132+
server.injectWorkerMetadata(
4133+
WorkerMetadataResponse.newBuilder()
4134+
.setMetadataVersion(1)
4135+
.addWorkEndpoints(
4136+
WorkerMetadataResponse.Endpoint.newBuilder()
4137+
.setBackendWorkerToken("workerToken1")
4138+
.setDirectEndpoint(
4139+
WindmillServiceAddress.create(HostAndPort.fromParts("localhost", 12345))
4140+
.gcpServiceAddress()
4141+
.toString())
4142+
.build())
4143+
.build());
4144+
41334145
// Switch to Directpath.
41344146
StreamingGlobalConfig directPathConfig =
41354147
StreamingGlobalConfig.builder()
@@ -4144,10 +4156,10 @@ public void testSwitchStreamingWorkerHarness() throws Exception {
41444156
}
41454157

41464158
// Wait for the harnessSwitchExecutor to complete the switch.
4147-
Future<?> future = harnessSwitchExecutor.submit(() -> {});
4159+
Future<?> directPathSwitchFuture = harnessSwitchExecutor.submit(() -> {});
41484160
// Wait for the dummy task to complete. The dummy task will be executed after
41494161
// switchStreamingWorkerHarness has completed.
4150-
future.get(15, TimeUnit.SECONDS);
4162+
directPathSwitchFuture.get(30, TimeUnit.SECONDS);
41514163
assertTrue(
41524164
"Worker should switch to FanOutStreamingEngineWorkerHarness",
41534165
worker.getStreamingWorkerHarness() instanceof FanOutStreamingEngineWorkerHarness);
@@ -4164,8 +4176,8 @@ public void testSwitchStreamingWorkerHarness() throws Exception {
41644176
observer.accept(cloudPathConfig);
41654177
}
41664178

4167-
future = harnessSwitchExecutor.submit(() -> {});
4168-
future.get(15, TimeUnit.SECONDS);
4179+
Future<?> cloudPathSwitchFuture = harnessSwitchExecutor.submit(() -> {});
4180+
cloudPathSwitchFuture.get(30, TimeUnit.SECONDS);
41694181
assertTrue(
41704182
"Worker should switch back to SingleSourceWorkerHarness",
41714183
worker.getStreamingWorkerHarness() instanceof SingleSourceWorkerHarness);

0 commit comments

Comments
 (0)