<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.scheduler.adapter;

import org.apache.flink.runtime.executiongraph.ExecutionEdge;
import org.apache.flink.runtime.executiongraph.ExecutionGraph;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.runtime.executiongraph.IntermediateResultPartition;
import org.apache.flink.runtime.executiongraph.TestRestartStrategy;
import org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGateway;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.apache.flink.runtime.scheduler.strategy.ResultPartitionState;
import org.apache.flink.util.IterableUtils;
import org.apache.flink.util.TestLogger;

import org.apache.flink.shaded.guava18.com.google.common.collect.Iterables;

import org.junit.Before;
import org.junit.Test;

import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static junit.framework.TestCase.assertTrue;
import static org.apache.flink.api.common.InputDependencyConstraint.ALL;
import static org.apache.flink.api.common.InputDependencyConstraint.ANY;
import static org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.createNoOpVertex;
import static org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.createSimpleTestGraph;
import static org.apache.flink.runtime.io.network.partition.ResultPartitionType.PIPELINED;
import static org.apache.flink.runtime.jobgraph.DistributionPattern.ALL_TO_ALL;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.fail;

/**
 * Unit tests for {@link DefaultExecutionTopology}.
 */
public class DefaultExecutionTopologyTest extends TestLogger {

	private final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();

	private final TestRestartStrategy triggeredRestartStrategy = TestRestartStrategy.manuallyTriggered();

	private ExecutionGraph executionGraph;

	private DefaultExecutionTopology adapter;

	@Before
	public void setUp() throws Exception {
		JobVertex[] jobVertices = new JobVertex[2];
		int parallelism = 3;
		jobVertices[0] = createNoOpVertex(parallelism);
		jobVertices[1] = createNoOpVertex(parallelism);
		jobVertices[1].connectNewDataSetAsInput(jobVertices[0], ALL_TO_ALL, PIPELINED);
		jobVertices[0].setInputDependencyConstraint(ALL);
		jobVertices[1].setInputDependencyConstraint(ANY);
		executionGraph = createSimpleTestGraph(
			taskManagerGateway,
			triggeredRestartStrategy,
			jobVertices);
		adapter = new DefaultExecutionTopology(executionGraph);
	}

	@Test
	public void testConstructor() {
		// implicitly tests order constraint of getVertices()
		assertGraphEquals(executionGraph, adapter);
	}

	@Test
	public void testGetResultPartition() {
		for (ExecutionVertex vertex : executionGraph.getAllExecutionVertices()) {
			for (Map.Entry<IntermediateResultPartitionID, IntermediateResultPartition> entry : vertex.getProducedPartitions().entrySet()) {
				IntermediateResultPartition partition = entry.getValue();
				DefaultResultPartition schedulingResultPartition = adapter.getResultPartition(entry.getKey())
					.orElseThrow(() -> new IllegalArgumentException("can not find partition " + entry.getKey()));

				assertPartitionEquals(partition, schedulingResultPartition);
			}
		}
	}

	@Test
	public void testResultPartitionStateSupplier() {
		final IntermediateResultPartition intermediateResultPartition = IterableUtils
			.toStream(executionGraph.getAllExecutionVertices())
			.flatMap(v -> v.getProducedPartitions().values().stream())
			.findAny()
			.get();

		final DefaultResultPartition schedulingResultPartition = adapter
			.getResultPartition(intermediateResultPartition.getPartitionId())
			.get();

		assertEquals(ResultPartitionState.CREATED, schedulingResultPartition.getState());

		intermediateResultPartition.markDataProduced();
		assertEquals(ResultPartitionState.CONSUMABLE, schedulingResultPartition.getState());
	}

	@Test
	public void testGetVertexOrThrow() {
		try {
			adapter.getVertexOrThrow(new ExecutionVertexID(new JobVertexID(), 0));<commit_msg>[FLINK-17050][runtime] Replace usages of getResultPartition() with getResultPartitionOrThrow()<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.scheduler.adapter;

import org.apache.flink.runtime.executiongraph.ExecutionEdge;
import org.apache.flink.runtime.executiongraph.ExecutionGraph;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.runtime.executiongraph.IntermediateResultPartition;
import org.apache.flink.runtime.executiongraph.TestRestartStrategy;
import org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGateway;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.apache.flink.runtime.scheduler.strategy.ResultPartitionState;
import org.apache.flink.util.IterableUtils;
import org.apache.flink.util.TestLogger;

import org.apache.flink.shaded.guava18.com.google.common.collect.Iterables;

import org.junit.Before;
import org.junit.Test;

import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static junit.framework.TestCase.assertTrue;
import static org.apache.flink.api.common.InputDependencyConstraint.ALL;
import static org.apache.flink.api.common.InputDependencyConstraint.ANY;
import static org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.createNoOpVertex;
import static org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.createSimpleTestGraph;
import static org.apache.flink.runtime.io.network.partition.ResultPartitionType.PIPELINED;
import static org.apache.flink.runtime.jobgraph.DistributionPattern.ALL_TO_ALL;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.fail;

/**
 * Unit tests for {@link DefaultExecutionTopology}.
 */
public class DefaultExecutionTopologyTest extends TestLogger {

	private final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();

	private final TestRestartStrategy triggeredRestartStrategy = TestRestartStrategy.manuallyTriggered();

	private ExecutionGraph executionGraph;

	private DefaultExecutionTopology adapter;

	@Before
	public void setUp() throws Exception {
		JobVertex[] jobVertices = new JobVertex[2];
		int parallelism = 3;
		jobVertices[0] = createNoOpVertex(parallelism);
		jobVertices[1] = createNoOpVertex(parallelism);
		jobVertices[1].connectNewDataSetAsInput(jobVertices[0], ALL_TO_ALL, PIPELINED);
		jobVertices[0].setInputDependencyConstraint(ALL);
		jobVertices[1].setInputDependencyConstraint(ANY);
		executionGraph = createSimpleTestGraph(
			taskManagerGateway,
			triggeredRestartStrategy,
			jobVertices);
		adapter = new DefaultExecutionTopology(executionGraph);
	}

	@Test
	public void testConstructor() {
		// implicitly tests order constraint of getVertices()
		assertGraphEquals(executionGraph, adapter);
	}

	@Test
	public void testGetResultPartition() {
		for (ExecutionVertex vertex : executionGraph.getAllExecutionVertices()) {
			for (Map.Entry<IntermediateResultPartitionID, IntermediateResultPartition> entry : vertex.getProducedPartitions().entrySet()) {
				IntermediateResultPartition partition = entry.getValue();
				DefaultResultPartition schedulingResultPartition = adapter.getResultPartitionOrThrow(entry.getKey());

				assertPartitionEquals(partition, schedulingResultPartition);
			}
		}
	}

	@Test
	public void testResultPartitionStateSupplier() {
		final IntermediateResultPartition intermediateResultPartition = IterableUtils
			.toStream(executionGraph.getAllExecutionVertices())
			.flatMap(v -> v.getProducedPartitions().values().stream())
			.findAny()
			.get();

		final DefaultResultPartition schedulingResultPartition = adapter
			.getResultPartitionOrThrow(intermediateResultPartition.getPartitionId());

		assertEquals(ResultPartitionState.CREATED, schedulingResultPartition.getState());

		intermediateResultPartition.markDataProduced();
		assertEquals(ResultPartitionState.CONSUMABLE, schedulingResultPartition.getState());
	}

	@Test
	public void testGetVertexOrThrow() {
		try {
			adapter.getVertexOrThrow(new ExecutionVertexID(new JobVertexID(), 0));
{'fraction_non_alphanumeric': 0.08806844601506557, 'fraction_numerical': 0.003254905607737376, 'mean_word_length': 6.280974949221394, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7665996', 'n_tokens_mistral': 3062, 'n_tokens_neox': 2896, 'n_words': 647}