Github user zentol commented on a diff in the pull request:
https://github.com/apache/flink/pull/6338#discussion_r202606214
--- Diff:
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/TestInputChannel.java
---
@@ -81,26 +72,21 @@ public TestInputChannel readBuffer(boolean
moreAvailable) throws IOException, In
}
public TestInputChannel readEndOfPartitionEvent() throws IOException,
InterruptedException {
- final Answer<Optional<BufferAndAvailability>> answer = new
Answer<Optional<BufferAndAvailability>>() {
- @Override
- public Optional<BufferAndAvailability>
answer(InvocationOnMock invocationOnMock) throws Throwable {
- // Return true after finishing
- when(mock.isReleased()).thenReturn(true);
-
- return Optional.of(new
BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
false, 0));
+ mock.addBufferAndAvailability(
+ new BufferAvailabilityProvider() {
+ @Override
+ public Optional<BufferAndAvailability>
getBufferAvailability() throws IOException, InterruptedException {
+ mock.setReleased();
+ return Optional.of(new
BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
+ false,
+ 0));
+ }
}
- };
-
- if (stubbing == null) {
- stubbing =
when(mock.getNextBuffer()).thenAnswer(answer);
- } else {
- stubbing = stubbing.thenAnswer(answer);
- }
-
+ );
return this;
}
- public InputChannel getInputChannel() {
+ public MockInputChannel getInputChannel() {
--- End diff --
It would be good if we could get by without exposing the mock. As far as i
can tell on GitHub the only usages of `MockInputChannel` methods outside of
this class are in `StreamTestSingleInputGate`:
```inputChannels[channelIndex].getInputChannel().addBufferAndAvailability(answer);```
```inputChannels[channelIndex].getInputChannel().setReleased();```
Since the `TestInputChannel` class is already accessed anyway we could move
these methods to the TestChannel class.
Note that currently this exposes a package-private class with a public
method, which means that anyone without package-private access will get a
compile error. Either make this method package private, or make the class
public.
---
