/*
 * 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.rpc.akka;

import org.apache.flink.api.common.time.Time;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.testutils.FlinkMatchers;
import org.apache.flink.runtime.rpc.RpcUtils;
import org.apache.flink.runtime.rpc.exceptions.RecipientUnreachableException;
import org.apache.flink.util.SerializedValue;
import org.apache.flink.util.TestLogger;

import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;

import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;

import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;

/** Tests for remote AkkaRpcActors. */
public class RemoteAkkaRpcActorTest extends TestLogger {

    private static AkkaRpcService rpcService;
    private static AkkaRpcService otherRpcService;

    private static final Configuration configuration = new Configuration();

    @BeforeClass
    public static void setupClass() throws Exception {
        rpcService =
                AkkaRpcServiceUtils.createRemoteRpcService(
                        configuration, "localhost", "0", null, Optional.empty());

        otherRpcService =
                AkkaRpcServiceUtils.createRemoteRpcService(
                        configuration, "localhost", "0", null, Optional.empty());
    }

    @AfterClass
    public static void teardownClass()
            throws InterruptedException, ExecutionException, TimeoutException {
        RpcUtils.terminateRpcServices(Time.seconds(10), rpcService, otherRpcService);
    }

    @Test
    public void canRespondWithNullValueRemotely() throws Exception {
        try (final AkkaRpcActorTest.NullRespondingEndpoint nullRespondingEndpoint =
                new AkkaRpcActorTest.NullRespondingEndpoint(rpcService)) {
            nullRespondingEndpoint.start();

            final AkkaRpcActorTest.NullRespondingGateway rpcGateway =
                    otherRpcService
                            .connect(
                                    nullRespondingEndpoint.getAddress(),
                                    AkkaRpcActorTest.NullRespondingGateway.class)
                            .join();

            final CompletableFuture<Integer> nullValuedResponseFuture = rpcGateway.foobar();

            assertThat(nullValuedResponseFuture.join(), is(nullValue()));
        }
    }

    @Test
    public void canRespondWithSynchronousNullValueRemotely() throws Exception {
        try (final AkkaRpcActorTest.NullRespondingEndpoint nullRespondingEndpoint =
                new AkkaRpcActorTest.NullRespondingEndpoint(rpcService)) {
            nullRespondingEndpoint.start();

            final AkkaRpcActorTest.NullRespondingGateway rpcGateway =
                    otherRpcService
                            .connect(
                                    nullRespondingEndpoint.getAddress(),
                                    AkkaRpcActorTest.NullRespondingGateway.class)
                            .join();

            final Integer value = rpcGateway.synchronousFoobar();

            assertThat(value, is(nullValue()));
        }
    }

    @Test
    public void canRespondWithSerializedValueRemotely() throws Exception {
        try (final AkkaRpcActorTest.SerializedValueRespondingEndpoint endpoint =
                new AkkaRpcActorTest.SerializedValueRespondingEndpoint(rpcService)) {
            endpoint.start();

            final AkkaRpcActorTest.SerializedValueRespondingGateway remoteGateway =
                    otherRpcService
                            .connect(
                                    endpoint.getAddress(),
                                    AkkaRpcActorTest.SerializedValueRespondingGateway.class)
                            .join();

            assertThat(
                    remoteGateway.getSerializedValueSynchronously(),
                    equalTo(AkkaRpcActorTest.SerializedValueRespondingEndpoint.SERIALIZED_VALUE));

            final CompletableFuture<SerializedValue<String>> responseFuture =
                    remoteGateway.getSerializedValue();

            assertThat(
                    responseFuture.get(),
                    equalTo(AkkaRpcActorTest.SerializedValueRespondingEndpoint.SERIALIZED_VALUE));
        }
    }

    @Test
    public void failsRpcResultImmediatelyIfEndpointIsStopped() throws Exception {
        try (final AkkaRpcActorTest.SerializedValueRespondingEndpoint endpoint =
                new AkkaRpcActorTest.SerializedValueRespondingEndpoint(rpcService)) {
            endpoint.start();

            final AkkaRpcActorTest.SerializedValueRespondingGateway gateway =
                    otherRpcService
                            .connect(
                                    endpoint.getAddress(),
                                    AkkaRpcActorTest.SerializedValueRespondingGateway.class)
                            .join();

            endpoint.close();

            try {
                gateway.getSerializedValue().join();
                fail("The endpoint should have been stopped.");
            } catch (Exception e) {
                // the rpc result should not fail because of a TimeoutException
                assertThat(e, FlinkMatchers.containsCause(RecipientUnreachableException.class));
            }
        }
    }

    @Test
    public void failsRpcResultImmediatelyIfRemoteRpcServiceIsNotAvailable() throws Exception {
        final AkkaRpcService toBeClosedRpcService =
                AkkaRpcServiceUtils.createRemoteRpcService(
                        configuration, "localhost", "0", null, Optional.empty());
        try (final AkkaRpcActorTest.SerializedValueRespondingEndpoint endpoint =
                new AkkaRpcActorTest.SerializedValueRespondingEndpoint(toBeClosedRpcService)) {
            endpoint.start();

            final AkkaRpcActorTest.SerializedValueRespondingGateway gateway =
                    otherRpcService
                            .connect(
                                    endpoint.getAddress(),
                                    AkkaRpcActorTest.SerializedValueRespondingGateway.class)
                            .join();

            toBeClosedRpcService.stopService().join();

            Thread.sleep(100L);

            try {
                gateway.getSerializedValue().join();
                fail("The endpoint should have been stopped.");
            } catch (Exception e) {
                // the rpc result should not fail because of a TimeoutException
                assertThat(e, FlinkMatchers.containsCause(RecipientUnreachableException.class));
            }
        } finally {
            RpcUtils.terminateRpcService(toBeClosedRpcService, Time.seconds(10L));
        }
    }
}

@@ -1460,16 +1460,69 @@
 ption;%0A%0A
+import static org.hamcrest.MatcherAssert.assertThat;%0A
 import s
@@ -1649,51 +1649,8 @@
 ue;%0A
-import static org.junit.Assert.assertThat;%0A
 impo

Replace junit.Assert.assertThat with hamcrest.MatcherAssert.assertThat in RemoteAkkaRpcServiceTest