flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/TwoInputStreamTaskTestHarness.java

/**
 * 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.streaming.runtime.tasks;

import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.io.network.partition.consumer.StreamTestSingleInputGate;
import org.apache.flink.streaming.api.collector.selector.OutputSelector;
import org.apache.flink.streaming.api.graph.StreamEdge;
import org.apache.flink.streaming.api.graph.StreamNode;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.StreamOperator;
import org.apache.flink.streaming.runtime.partitioner.BroadcastPartitioner;

import java.io.IOException;
import java.util.LinkedList;
import java.util.List;


/**
 * Test harness for testing a {@link TwoInputStreamTask}.
 *
 * <p>
 * This mock Invokable provides the task with a basic runtime context and allows pushing elements
 * and watermarks into the task. {@link #getOutput()} can be used to get the emitted elements
 * and events. You are free to modify the retrieved list.
 *
 * <p>
 * After setting up everything the Task can be invoked using {@link #invoke()}. This will start
 * a new Thread to execute the Task. Use {@link #waitForTaskCompletion()} to wait for the Task
 * thread to finish. Use {@link #processElement}
 * to send elements to the task. Use
 * {@link #processEvent(org.apache.flink.runtime.event.AbstractEvent)} to send events to the task.
 * Before waiting for the task to finish you must call {@link #endInput()} to signal to the task
 * that data entry is finished.
 *
 * <p>
 * When Elements or Events are offered to the Task they are put into a queue. The input gates
 * of the Task read from this queue. Use {@link #waitForInputProcessing()} to wait until all
 * queues are empty. This must be used after entering some elements before checking the
 * desired output.
 *
 * <p>
 * When using this you need to add the following line to your test class to setup Powermock:
 * {@code {@literal @}PrepareForTest({ResultPartitionWriter.class})}
 */
public class TwoInputStreamTaskTestHarness<IN1, IN2, OUT> extends StreamTaskTestHarness<OUT> {

	private TypeInformation<IN1> inputType1;
	private TypeSerializer<IN1> inputSerializer1;

	private TypeInformation<IN2> inputType2;
	private TypeSerializer<IN2> inputSerializer2;

	private int[] inputGateAssignment;

	/**
	 * Creates a test harness with the specified number of input gates and specified number
	 * of channels per input gate. Parameter inputGateAssignment specifies for each gate whether
	 * it should be assigned to the first (1), or second (2) input of the task.
	 */
	public TwoInputStreamTaskTestHarness(TwoInputStreamTask<IN1, IN2, OUT> task,
			int numInputGates,
			int numInputChannelsPerGate,
			int[] inputGateAssignment,
			TypeInformation<IN1> inputType1,
			TypeInformation<IN2> inputType2,
			TypeInformation<OUT> outputType) {
		super(task, outputType);

		this.inputType1 = inputType1;
		inputSerializer1 = inputType1.createSerializer(executionConfig);

		this.inputType2 = inputType2;
		inputSerializer2 = inputType2.createSerializer(executionConfig);

		this.numInputGates = numInputGates;
		this.numInputChannelsPerGate = numInputChannelsPerGate;
		this.inputGateAssignment = inputGateAssignment;
	}

	/**
	 * Creates a test harness with one input gate (that has one input channel) per input. The first
	 * input gate is assigned to the first task input, the second input gate is assigned to the
	 * second task input.
	 */
	public TwoInputStreamTaskTestHarness(TwoInputStreamTask<IN1, IN2, OUT> task,
			TypeInformation<IN1> inputType1,
			TypeInformation<IN2> inputType2,
			TypeInformation<OUT> outputType) {
		this(task, 2, 1, new int[] {1, 2}, inputType1, inputType2, outputType);
	}

	@Override
	protected void initializeInputs() throws IOException, InterruptedException {

		inputGates = new StreamTestSingleInputGate[numInputGates];
		List<StreamEdge> inPhysicalEdges = new LinkedList<StreamEdge>();

		StreamOperator<IN1> dummyOperator = new AbstractStreamOperator<IN1>() {
			private static final long serialVersionUID = 1L;
		};

		StreamNode sourceVertexDummy = new StreamNode(null, 0, "default group", dummyOperator, "source dummy", new LinkedList<OutputSelector<?>>(), SourceStreamTask.class);
		StreamNode targetVertexDummy = new StreamNode(null, 1, "default group", dummyOperator, "target dummy", new LinkedList<OutputSelector<?>>(), SourceStreamTask.class);

		for (int i = 0; i < numInputGates; i++) {

			switch (inputGateAssignment[i]) {
				case 1: {
					inputGates[i] = new StreamTestSingleInputGate<IN1>(
							numInputChannelsPerGate,
							bufferSize,
							inputSerializer1);


					StreamEdge streamEdge = new StreamEdge(sourceVertexDummy,
							targetVertexDummy,
							1,
							new LinkedList<String>(),
							new BroadcastPartitioner<Object>());

					inPhysicalEdges.add(streamEdge);
					break;
				}
				case 2: {
					inputGates[i] = new StreamTestSingleInputGate<IN2>(
							numInputChannelsPerGate,
							bufferSize,
							inputSerializer2);

					StreamEdge streamEdge = new StreamEdge(sourceVertexDummy,
							targetVertexDummy,
							2,
							new LinkedList<String>(),
							new BroadcastPartitioner<Object>());

					inPhysicalEdges.add(streamEdge);
					break;
				}
				default:
					throw new IllegalStateException("Wrong input gate assignment.");
			}

			this.mockEnv.addInputGate(inputGates[i].getInputGate());
		}

		streamConfig.setInPhysicalEdges(inPhysicalEdges);
		streamConfig.setNumberOfInputs(numInputGates);
		streamConfig.setTypeSerializerIn1(inputSerializer1);
		streamConfig.setTypeSerializerIn2(inputSerializer2);
	}

}