@Override
    public InputChannel getChannel(int channelIndex) {
        int gateIndex = inputChannelToInputGateIndex[channelIndex];
        return inputGatesByGateIndex
                .get(gateIndex)
                .getChannel(channelIndex - inputGateChannelIndexOffsets[gateIndex]);
    }
@Test
    void testUpdateInputChannel() throws Exception {
        final SingleInputGate inputGate1 = createInputGate(1);
        TestInputChannel inputChannel1 = new TestInputChannel(inputGate1, 0);
        inputGate1.setInputChannels(inputChannel1);

        final SingleInputGate inputGate2 = createInputGate(1);
        TestingResultPartitionManager partitionManager =
                new TestingResultPartitionManager(new NoOpResultSubpartitionView());
        InputChannel unknownInputChannel2 =
                InputChannelBuilder.newBuilder()
                        .setPartitionManager(partitionManager)
                        .buildUnknownChannel(inputGate2);
        inputGate2.setInputChannels(unknownInputChannel2);

        UnionInputGate unionInputGate = new UnionInputGate(inputGate1, inputGate2);
        ResultPartitionID resultPartitionID = unknownInputChannel2.getPartitionId();
        ResourceID location = ResourceID.generate();
        inputGate2.updateInputChannel(
                location,
                createRemoteWithIdAndLocation(resultPartitionID.getPartitionId(), location));

        assertThat(unionInputGate.getChannel(0)).isEqualTo(inputChannel1);
        // Check that updated input channel is visible via UnionInputGate
        assertThat(unionInputGate.getChannel(1)).isEqualTo(inputGate2.getChannel(0));
    }