content: /*
   * 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.celeborn.plugin.flink;
  
  import static org.mockito.ArgumentMatchers.*;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  import java.io.IOException;
  import java.nio.ByteBuffer;
  
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.Unpooled;
  import io.netty.channel.ChannelFuture;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.network.client.RpcResponseCallback;
  import org.apache.celeborn.common.network.client.TransportClient;
  import org.apache.celeborn.common.network.client.TransportClientFactory;
  import org.apache.celeborn.common.protocol.PartitionLocation;
  import org.apache.celeborn.common.protocol.message.StatusCode;
  import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
  
  public class FlinkShuffleClientImplSuiteJ {
    static int BufferSize = 64;
    static byte[] TEST_BUF1 = new byte[BufferSize];
    protected ChannelFuture mockedFuture = mock(ChannelFuture.class);
    static CelebornConf conf;
    static FlinkShuffleClientImpl shuffleClient;
    protected static final TransportClientFactory clientFactory = mock(TransportClientFactory.class);
    protected final TransportClient client = mock(TransportClient.class);
    protected static final PartitionLocation primaryLocation =
        new PartitionLocation(0, 1, "localhost", 1, 1, 1, 1, PartitionLocation.Mode.PRIMARY);
  
    @Before
    public void setup() throws IOException, InterruptedException {
      conf = new CelebornConf();
      shuffleClient =
          new FlinkShuffleClientImpl(
              "APP", "localhost", 1232, System.currentTimeMillis(), conf, null, -1) {
            @Override
            public void setupLifecycleManagerRef(String host, int port) {}
          };
      when(clientFactory.createClient(primaryLocation.getHost(), primaryLocation.getPushPort(), 1))
          .thenAnswer(t -> client);
  
      shuffleClient.setDataClientFactory(clientFactory);
    }
  
    public ByteBuf createByteBuf() {
      for (int i = 16; i < BufferSize; i++) {
        TEST_BUF1[i] = 1;
      }
      ByteBuf byteBuf = Unpooled.wrappedBuffer(TEST_BUF1);
      byteBuf.writerIndex(BufferSize);
      return byteBuf;
    }
  
    @Test
    public void testPushDataByteBufSuccess() throws IOException {
      ByteBuf byteBuf = createByteBuf();
      when(client.pushData(any(), anyLong(), any()))
          .thenAnswer(
              t -> {
                RpcResponseCallback rpcResponseCallback = t.getArgument(1, RpcResponseCallback.class);
                ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[0]);
                rpcResponseCallback.onSuccess(byteBuffer);
                return mockedFuture;
              });
  
      int pushDataLen =
          shuffleClient.pushDataToLocation(2, 3, 4, 5, byteBuf, primaryLocation, () -> {});
      Assert.assertEquals(BufferSize, pushDataLen);
    }
  
    @Test
    public void testPushDataByteBufHardSplit() throws IOException {
      ByteBuf byteBuf = Unpooled.wrappedBuffer(TEST_BUF1);
      when(client.pushData(any(), anyLong(), any()))
          .thenAnswer(
              t -> {
                RpcResponseCallback rpcResponseCallback = t.getArgument(1, RpcResponseCallback.class);
                ByteBuffer byteBuffer =
                    ByteBuffer.wrap(new byte[] {StatusCode.HARD_SPLIT.getValue()});
                rpcResponseCallback.onSuccess(byteBuffer);
                return mockedFuture;
              });
      int pushDataLen =
          shuffleClient.pushDataToLocation(2, 3, 4, 5, byteBuf, primaryLocation, () -> {});
      Assert.assertEquals(BufferSize, pushDataLen);
    }
  
    @Test
    public void testPushDataByteBufFail() throws IOException {
      ByteBuf byteBuf = Unpooled.wrappedBuffer(TEST_BUF1);
      when(client.pushData(any(), anyLong(), any(), any()))
          .thenAnswer(
              t -> {
                RpcResponseCallback rpcResponseCallback = t.getArgument(1, RpcResponseCallback.class);
                rpcResponseCallback.onFailure(new Exception("pushDataFailed"));
                return mockedFuture;
              });
      // first push just  set pushdata.exception
      shuffleClient.pushDataToLocation(2, 3, 4, 5, byteBuf, primaryLocation, () -> {});
  
      boolean isFailed = false;
      // second push will throw exception
      try {
        shuffleClient.pushDataToLocation(2, 3, 4, 5, byteBuf, primaryLocation, () -> {});
      } catch (IOException e) {
        isFailed = true;
      } finally {
        Assert.assertTrue("should failed", isFailed);
      }
    }
  }
  