From aa8a2ba894bc64ac3ce492ea3fecf427a2cf9607 Mon Sep 17 00:00:00 2001 From: Chi Wang Date: Wed, 28 Jun 2023 14:38:10 +0200 Subject: [PATCH] Fix remote execution tests --- .../build/lib/remote/GrpcCacheClientTest.java | 10 ++-- ...SpawnRunnerWithGrpcRemoteExecutorTest.java | 52 +++++++++---------- 2 files changed, 30 insertions(+), 32 deletions(-) diff --git a/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java b/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java index 26e374d05a958c..cf159a20e660b1 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java @@ -21,6 +21,7 @@ import static org.mockito.AdditionalAnswers.answerVoid; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; import build.bazel.remote.execution.v2.Action; @@ -811,10 +812,9 @@ public void updateActionResult( } } }); - ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class); + ByteStreamImplBase mockByteStreamImpl = spy(ByteStreamImplBase.class); serviceRegistry.addService(mockByteStreamImpl); - when(mockByteStreamImpl.write(ArgumentMatchers.>any())) - .thenAnswer( + doAnswer( new Answer>() { private int numErrors = 4; @@ -865,7 +865,9 @@ public void onError(Throwable t) { } }; } - }); + }) + .when(mockByteStreamImpl) + .write(any()); doAnswer( answerVoid( (QueryWriteStatusRequest request, diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java index c03e0a6d7395f9..dd8536553f1d3c 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java @@ -22,6 +22,7 @@ import static org.mockito.AdditionalAnswers.answerVoid; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -729,9 +730,8 @@ public void findMissingBlobs( }; serviceRegistry.addService(ServerInterceptors.intercept(cas, new RequestHeadersValidator())); - ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class); - when(mockByteStreamImpl.write(ArgumentMatchers.>any())) - .thenAnswer(blobWriteAnswer("xyz".getBytes(UTF_8))); + ByteStreamImplBase mockByteStreamImpl = spy(ByteStreamImplBase.class); + doAnswer(blobWriteAnswer("xyz".getBytes(UTF_8))).when(mockByteStreamImpl).write(any()); serviceRegistry.addService( ServerInterceptors.intercept(mockByteStreamImpl, new RequestHeadersValidator())); @@ -807,7 +807,7 @@ public void getActionResult( .setResponse(Any.pack(ExecuteResponse.newBuilder().setResult(actionResult).build())) .build(); - ExecutionImplBase mockExecutionImpl = Mockito.mock(ExecutionImplBase.class); + ExecutionImplBase mockExecutionImpl = spy(ExecutionImplBase.class); // Flow of this test: // - call execute, get retriable gRPC error // - retry: call execute, get retriable Operation error @@ -815,7 +815,7 @@ public void getActionResult( // - retry: call waitExecute, get a retriable gRPC error // - retry: call waitExecute, get retriable Operation error // - retry: call execute, get successful operation, ignore further errors. - Mockito.doAnswer(answerWith(null, Status.UNAVAILABLE)) + doAnswer(answerWith(null, Status.UNAVAILABLE)) .doAnswer(answerWith(operationWithExecuteError, Status.OK)) .doAnswer(answerWith(unfinishedOperation, Status.UNAVAILABLE)) .doAnswer(answerWith(opSuccess, Status.UNAVAILABLE)) // last status should be ignored. @@ -823,7 +823,7 @@ public void getActionResult( .execute( ArgumentMatchers.any(), ArgumentMatchers.>any()); - Mockito.doAnswer(answerWith(null, Status.UNAVAILABLE)) + doAnswer(answerWith(null, Status.UNAVAILABLE)) .doAnswer(answerWith(operationWithExecuteError, Status.OK)) .when(mockExecutionImpl) .waitExecution( @@ -853,11 +853,12 @@ public void findMissingBlobs( } }); - ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class); - when(mockByteStreamImpl.write(ArgumentMatchers.>any())) - .thenAnswer(blobWriteAnswerError()) // Error on the input file. - .thenAnswer(blobWriteAnswerError()) // Error on the input file again. - .thenAnswer(blobWriteAnswer("xyz".getBytes(UTF_8))); // Upload input file successfully. + ByteStreamImplBase mockByteStreamImpl = spy(ByteStreamImplBase.class); + doAnswer(blobWriteAnswerError()) // Error on the input file. + .doAnswer(blobWriteAnswerError()) // Error on the input file again. + .doAnswer(blobWriteAnswer("xyz".getBytes(UTF_8))) // Upload input file successfully. + .when(mockByteStreamImpl) + .write(any()); doAnswer( answerVoid( (QueryWriteStatusRequest request, @@ -871,7 +872,7 @@ public void findMissingBlobs( })) .when(mockByteStreamImpl) .queryWriteStatus(any(), any()); - Mockito.doAnswer( + doAnswer( invocationOnMock -> { @SuppressWarnings("unchecked") StreamObserver responseObserver = @@ -951,26 +952,22 @@ public void getActionResult( .setResponse(Any.pack(ExecuteResponse.newBuilder().setResult(actionResult).build())) .build(); - ExecutionImplBase mockExecutionImpl = Mockito.mock(ExecutionImplBase.class); + ExecutionImplBase mockExecutionImpl = spy(ExecutionImplBase.class); // Flow of this test: // - call execute, get an Operation, then a retriable gRPC error // - retry: call waitExecute, get NOT_FOUND (operation lost) // - retry: call execute, get NOT_FOUND (operation lost) // - retry: call execute, get an Operation, then a retriable gRPC error // - retry: call waitExecute, get successful operation, ignore further errors. - Mockito.doAnswer(answerWith(unfinishedOperation, Status.UNAVAILABLE)) + doAnswer(answerWith(unfinishedOperation, Status.UNAVAILABLE)) .doAnswer(answerWith(unfinishedOperation, Status.NOT_FOUND)) .doAnswer(answerWith(unfinishedOperation, Status.UNAVAILABLE)) .when(mockExecutionImpl) - .execute( - ArgumentMatchers.any(), - ArgumentMatchers.>any()); - Mockito.doAnswer(answerWith(unfinishedOperation, Status.NOT_FOUND)) + .execute(any(), any()); + doAnswer(answerWith(unfinishedOperation, Status.NOT_FOUND)) .doAnswer(answerWith(opSuccess, Status.UNAVAILABLE)) // This error is ignored. .when(mockExecutionImpl) - .waitExecution( - ArgumentMatchers.any(), - ArgumentMatchers.>any()); + .waitExecution(any(), any()); serviceRegistry.addService(mockExecutionImpl); serviceRegistry.addService( @@ -989,10 +986,9 @@ public void findMissingBlobs( } }); - ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class); - when(mockByteStreamImpl.write(ArgumentMatchers.>any())) - .thenAnswer(blobWriteAnswer("xyz".getBytes(UTF_8))); // Upload input file successfully. - Mockito.doAnswer( + ByteStreamImplBase mockByteStreamImpl = spy(ByteStreamImplBase.class); + doAnswer(blobWriteAnswer("xyz".getBytes(UTF_8))).when(mockByteStreamImpl).write(any()); + doAnswer( invocationOnMock -> { @SuppressWarnings("unchecked") StreamObserver responseObserver = @@ -1508,17 +1504,17 @@ public void read(ReadRequest request, StreamObserver responseObser .build(); final WaitExecutionRequest waitExecutionRequest = WaitExecutionRequest.newBuilder().setName(opName).build(); - ExecutionImplBase mockExecutionImpl = Mockito.mock(ExecutionImplBase.class); + ExecutionImplBase mockExecutionImpl = spy(ExecutionImplBase.class); // Flow of this test: // - call execute, get an unfinished Operation, then the stream completes // - call waitExecute, get an unfinished Operation, then the stream completes // - call waitExecute, get a finished Operation - Mockito.doAnswer(answerWith(unfinishedOperation, Status.OK)) + doAnswer(answerWith(unfinishedOperation, Status.OK)) .when(mockExecutionImpl) .execute( ArgumentMatchers.any(), ArgumentMatchers.>any()); - Mockito.doAnswer(answerWith(unfinishedOperation, Status.OK)) + doAnswer(answerWith(unfinishedOperation, Status.OK)) .doAnswer(answerWith(completeOperation, Status.OK)) .when(mockExecutionImpl) .waitExecution(