@Override
    public CompletableFuture<ExecutionState> requestPartitionState(
            final IntermediateDataSetID intermediateResultId,
            final ResultPartitionID resultPartitionId) {

        try {
            return CompletableFuture.completedFuture(
                    schedulerNG.requestPartitionState(intermediateResultId, resultPartitionId));
        } catch (PartitionProducerDisposedException e) {
            log.info("Error while requesting partition state", e);
            return FutureUtils.completedExceptionally(e);
        }
    }
@Test
    void testRequestPartitionState() throws Exception {
        final JobGraph producerConsumerJobGraph = producerConsumerJobGraph();
        try (final JobMaster jobMaster =
                new JobMasterBuilder(producerConsumerJobGraph, rpcService)
                        .withConfiguration(configuration)
                        .withHighAvailabilityServices(haServices)
                        .withHeartbeatServices(heartbeatServices)
                        .createJobMaster()) {

            jobMaster.start();

            final CompletableFuture<TaskDeploymentDescriptor> tddFuture = new CompletableFuture<>();
            final TestingTaskExecutorGateway testingTaskExecutorGateway =
                    new TestingTaskExecutorGatewayBuilder()
                            .setSubmitTaskConsumer(
                                    (taskDeploymentDescriptor, jobMasterId) -> {
                                        tddFuture.complete(taskDeploymentDescriptor);
                                        return CompletableFuture.completedFuture(Acknowledge.get());
                                    })
                            .createTestingTaskExecutorGateway();
            final LocalUnresolvedTaskManagerLocation taskManagerLocation =
                    new LocalUnresolvedTaskManagerLocation();

            final JobMasterGateway jobMasterGateway =
                    jobMaster.getSelfGateway(JobMasterGateway.class);

            final Collection<SlotOffer> slotOffers =
                    registerSlotsAtJobMaster(
                            1,
                            jobMasterGateway,
                            producerConsumerJobGraph.getJobID(),
                            testingTaskExecutorGateway,
                            taskManagerLocation);

            assertThat(slotOffers).hasSize(1);

            // obtain tdd for the result partition ids
            final TaskDeploymentDescriptor tdd = tddFuture.get();

            assertThat(tdd.getProducedPartitions()).hasSize(1);
            final ResultPartitionDeploymentDescriptor partition =
                    tdd.getProducedPartitions().iterator().next();

            final ExecutionAttemptID executionAttemptId = tdd.getExecutionAttemptId();
            final ExecutionAttemptID copiedExecutionAttemptId =
                    InstantiationUtil.clone(executionAttemptId);

            // finish the producer task
            jobMasterGateway
                    .updateTaskExecutionState(
                            SchedulerTestingUtils.createFinishedTaskExecutionState(
                                    executionAttemptId))
                    .get();

            // request the state of the result partition of the producer
            final ResultPartitionID partitionId =
                    new ResultPartitionID(partition.getPartitionId(), copiedExecutionAttemptId);
            CompletableFuture<ExecutionState> partitionStateFuture =
                    jobMasterGateway.requestPartitionState(partition.getResultId(), partitionId);

            assertThat(partitionStateFuture.get()).isEqualTo(ExecutionState.FINISHED);

            // ask for unknown result partition
            partitionStateFuture =
                    jobMasterGateway.requestPartitionState(
                            partition.getResultId(), new ResultPartitionID());

            assertThatThrownBy(partitionStateFuture::get)
                    .hasRootCauseInstanceOf(IllegalArgumentException.class);

            // ask for wrong intermediate data set id
            partitionStateFuture =
                    jobMasterGateway.requestPartitionState(
                            new IntermediateDataSetID(), partitionId);

            assertThatThrownBy(partitionStateFuture::get)
                    .hasRootCauseInstanceOf(IllegalArgumentException.class);

            // ask for "old" execution
            partitionStateFuture =
                    jobMasterGateway.requestPartitionState(
                            partition.getResultId(),
                            new ResultPartitionID(
                                    partition.getPartitionId(), createExecutionAttemptId()));

            assertThatThrownBy(partitionStateFuture::get)
                    .hasRootCauseInstanceOf(PartitionProducerDisposedException.class);
        }
    }