/** * Licensed 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.connectors.pulsar; import org.apache.flink.api.common.JobID; import org.apache.flink.client.program.ClusterClient; import org.apache.flink.runtime.client.JobStatusMessage; import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; import org.apache.flink.test.util.MiniClusterWithClientResource; import org.apache.pulsar.common.naming.TopicName; import org.junit.Before; import org.junit.ClassRule; import java.util.Collection; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; /** * Setup flink mini cluster. */ public abstract class PulsarTestBaseWithFlink extends PulsarTestBase { protected static final int NUM_TMS = 1; protected static final int TM_SLOTS = 8; protected ClusterClient<?> client; @ClassRule public static MiniClusterWithClientResource flink = new MiniClusterWithClientResource( new MiniClusterResourceConfiguration.Builder() .setConfiguration(getFlinkConfiguration()) .setNumberTaskManagers(NUM_TMS) .setNumberSlotsPerTaskManager(TM_SLOTS) .build()); @Before public void noJobIsRunning() throws Exception { client = flink.getClusterClient(); waitUntilNoJobIsRunning(client); } public static void waitUntilJobIsRunning(ClusterClient<?> client) throws Exception { while (getRunningJobs(client).isEmpty()) { Thread.sleep(50); } } public static void waitUntilNoJobIsRunning(ClusterClient<?> client) throws Exception { while (!getRunningJobs(client).isEmpty()) { Thread.sleep(50); cancelRunningJobs(client); } } public static List<JobID> getRunningJobs(ClusterClient<?> client) throws Exception { Collection<JobStatusMessage> statusMessages = client.listJobs().get(); return statusMessages.stream() .filter(status -> !status.getJobState().isGloballyTerminalState()) .map(JobStatusMessage::getJobId) .collect(Collectors.toList()); } public static void cancelRunningJobs(ClusterClient<?> client) throws Exception { List<JobID> runningJobs = getRunningJobs(client); for (JobID runningJob : runningJobs) { client.cancel(runningJob); } } public static final AtomicInteger TOPIC_ID = new AtomicInteger(0); public static String newTopic() { synchronized (TOPIC_ID) { int i = TOPIC_ID.getAndIncrement(); return TopicName.get("topic-" + i).toString(); } } }