/*
 * 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.test.util;

import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.client.program.MiniClusterClient;
import org.apache.flink.runtime.testutils.MiniClusterResource;
import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
import org.apache.flink.streaming.util.TestStreamEnvironment;

/**
 * Starts a Flink mini cluster as a resource and registers the respective ExecutionEnvironment and
 * StreamExecutionEnvironment.
 */
public class MiniClusterWithClientResource extends MiniClusterResource {

    private ClusterClient<?> clusterClient;

    private TestEnvironment executionEnvironment;

    public MiniClusterWithClientResource(
            final MiniClusterResourceConfiguration miniClusterResourceConfiguration) {
        super(miniClusterResourceConfiguration);
    }

    public ClusterClient<?> getClusterClient() {
        return clusterClient;
    }

    public TestEnvironment getTestEnvironment() {
        return executionEnvironment;
    }

    @Override
    public void before() throws Exception {
        super.before();

        clusterClient = createMiniClusterClient();

        executionEnvironment = new TestEnvironment(getMiniCluster(), getNumberSlots(), false);
        executionEnvironment.setAsContext();
        TestStreamEnvironment.setAsContext(getMiniCluster(), getNumberSlots());
    }

    @Override
    public void after() {
        log.info("Finalization triggered: Cluster shutdown is going to be initiated.");
        TestStreamEnvironment.unsetAsContext();
        TestEnvironment.unsetAsContext();

        Exception exception = null;

        if (clusterClient != null) {
            try {
                clusterClient.close();
            } catch (Exception e) {
                exception = e;
            }
        }

        clusterClient = null;

        super.after();

        if (exception != null) {
            log.warn("Could not properly shut down the MiniClusterWithClientResource.", exception);
        }
    }

    private MiniClusterClient createMiniClusterClient() {
        return new MiniClusterClient(getClientConfiguration(), getMiniCluster());
    }
}

@@ -941,32 +941,95 @@
 iClusterClient;%0A
+import org.apache.flink.client.program.rest.RestClusterClient;%0A
 import org.apach
@@ -1212,16 +1212,61 @@
 ronment;
+%0Aimport org.apache.flink.util.ExceptionUtils;
 %0A%0A/**%0A *
@@ -1510,24 +1510,106 @@
 usterClient;
+%0A    private RestClusterClient%3CMiniClusterClient.MiniClusterId%3E restClusterClient;
 %0A%0A    privat
@@ -1912,32 +1912,374 @@
 rClient;%0A    %7D%0A%0A
+    /**%0A     * Returns a %7B@link RestClusterClient%7D that can be used to communicate with this mini cluster.%0A     * Only use this if the client returned via %7B@link #getClusterClient()%7D does not fulfill your%0A     * needs.%0A     */%0A    public RestClusterClient%3C?%3E getRestClusterClient() throws Exception %7B%0A        return restClusterClient;%0A    %7D%0A%0A
     public TestE
@@ -2485,24 +2485,79 @@
 terClient();
+%0A        restClusterClient = createRestClusterClient();
 %0A%0A        ex
@@ -3238,32 +3238,306 @@
 Client = null;%0A%0A
+        if (restClusterClient != null) %7B%0A            try %7B%0A                restClusterClient.close();%0A            %7D catch (Exception e) %7B%0A                exception = ExceptionUtils.firstOrSuppressed(e, exception);%0A            %7D%0A        %7D%0A%0A        restClusterClient = null;%0A%0A
         super.af
@@ -3839,11 +3839,263 @@
 );%0A    %7D
+%0A%0A    private RestClusterClient%3CMiniClusterClient.MiniClusterId%3E createRestClusterClient()%0A            throws Exception %7B%0A        return new RestClusterClient%3C%3E(%0A                getClientConfiguration(), MiniClusterClient.MiniClusterId.INSTANCE);%0A    %7D
 %0A%7D%0A

Add RestClusterClient to MiniClusterWithClientResource