Skip to content

Commit

Permalink
Fix remote execution tests
Browse files Browse the repository at this point in the history
by replacing `mock` with `spy` because we rely on the real method and these methods are final.

Before mockito 5, it uses sublcass mockmaker which cannot mock final method/class. So when we call final methods on these mocked classes, we are calling into the real methods.

Now, it uses inline mockmaker which **can** mock final method/class. Calling final methods on these mocked classes will return `null` if we didn't provide the stub.

This change fixes that by using `spy`.

Closes #18800.

PiperOrigin-RevId: 544038459
Change-Id: Id75a5cb3b9d14e38d6bec918449e5aee671471eb
  • Loading branch information
coeuvre authored and copybara-github committed Jun 28, 2023
1 parent d055c46 commit 5637083
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 33 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +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.when;
import static org.mockito.Mockito.spy;

import build.bazel.remote.execution.v2.Action;
import build.bazel.remote.execution.v2.ActionCacheGrpc.ActionCacheImplBase;
Expand Down Expand Up @@ -811,10 +811,9 @@ public void updateActionResult(
}
}
});
ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class);
ByteStreamImplBase mockByteStreamImpl = spy(ByteStreamImplBase.class);
serviceRegistry.addService(mockByteStreamImpl);
when(mockByteStreamImpl.write(ArgumentMatchers.<StreamObserver<WriteResponse>>any()))
.thenAnswer(
doAnswer(
new Answer<StreamObserver<WriteRequest>>() {
private int numErrors = 4;

Expand Down Expand Up @@ -865,7 +864,9 @@ public void onError(Throwable t) {
}
};
}
});
})
.when(mockByteStreamImpl)
.write(any());
doAnswer(
answerVoid(
(QueryWriteStatusRequest request,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -730,9 +731,8 @@ public void findMissingBlobs(
};
serviceRegistry.addService(ServerInterceptors.intercept(cas, new RequestHeadersValidator()));

ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class);
when(mockByteStreamImpl.write(ArgumentMatchers.<StreamObserver<WriteResponse>>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()));

Expand Down Expand Up @@ -808,23 +808,23 @@ 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
// - retry: call execute, get an Operation, then a retriable gRPC error
// - 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.
.when(mockExecutionImpl)
.execute(
ArgumentMatchers.<ExecuteRequest>any(),
ArgumentMatchers.<StreamObserver<Operation>>any());
Mockito.doAnswer(answerWith(null, Status.UNAVAILABLE))
doAnswer(answerWith(null, Status.UNAVAILABLE))
.doAnswer(answerWith(operationWithExecuteError, Status.OK))
.when(mockExecutionImpl)
.waitExecution(
Expand Down Expand Up @@ -854,11 +854,12 @@ public void findMissingBlobs(
}
});

ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class);
when(mockByteStreamImpl.write(ArgumentMatchers.<StreamObserver<WriteResponse>>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,
Expand All @@ -872,7 +873,7 @@ public void findMissingBlobs(
}))
.when(mockByteStreamImpl)
.queryWriteStatus(any(), any());
Mockito.doAnswer(
doAnswer(
invocationOnMock -> {
@SuppressWarnings("unchecked")
StreamObserver<ReadResponse> responseObserver =
Expand Down Expand Up @@ -952,26 +953,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.<ExecuteRequest>any(),
ArgumentMatchers.<StreamObserver<Operation>>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.<WaitExecutionRequest>any(),
ArgumentMatchers.<StreamObserver<Operation>>any());
.waitExecution(any(), any());
serviceRegistry.addService(mockExecutionImpl);

serviceRegistry.addService(
Expand All @@ -990,10 +987,9 @@ public void findMissingBlobs(
}
});

ByteStreamImplBase mockByteStreamImpl = Mockito.mock(ByteStreamImplBase.class);
when(mockByteStreamImpl.write(ArgumentMatchers.<StreamObserver<WriteResponse>>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<ReadResponse> responseObserver =
Expand Down Expand Up @@ -1514,17 +1510,17 @@ public void read(ReadRequest request, StreamObserver<ReadResponse> 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.<ExecuteRequest>any(),
ArgumentMatchers.<StreamObserver<Operation>>any());
Mockito.doAnswer(answerWith(unfinishedOperation, Status.OK))
doAnswer(answerWith(unfinishedOperation, Status.OK))
.doAnswer(answerWith(completeOperation, Status.OK))
.when(mockExecutionImpl)
.waitExecution(
Expand Down

0 comments on commit 5637083

Please sign in to comment.