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.client;
  
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  import java.io.IOException;
  
  import scala.reflect.ClassTag$;
  
  import io.netty.channel.ChannelFuture;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.identity.UserIdentifier;
  import org.apache.celeborn.common.network.client.TransportClient;
  import org.apache.celeborn.common.network.client.TransportClientFactory;
  import org.apache.celeborn.common.protocol.CompressionCodec;
  import org.apache.celeborn.common.protocol.PartitionLocation;
  import org.apache.celeborn.common.protocol.PbRegisterShuffleResponse;
  import org.apache.celeborn.common.protocol.message.ControlMessages;
  import org.apache.celeborn.common.protocol.message.StatusCode;
  import org.apache.celeborn.common.rpc.RpcEndpointRef;
  
  public abstract class ShuffleClientBaseSuiteJ {
    protected ShuffleClientImpl shuffleClient = null;
    protected static final RpcEndpointRef endpointRef = mock(RpcEndpointRef.class);
    protected static final TransportClientFactory clientFactory = mock(TransportClientFactory.class);
    protected final TransportClient client = mock(TransportClient.class);
  
    protected static final String TEST_APPLICATION_ID = "testapp1";
    protected static final int TEST_SHUFFLE_ID = 1;
    protected static final int TEST_ATTEMPT_ID = 0;
    protected static final int TEST_REDUCRE_ID = 0;
  
    protected static final int PRIMARY_RPC_PORT = 1234;
    protected static final int PRIMARY_PUSH_PORT = 1235;
    protected static final int PRIMARY_FETCH_PORT = 1236;
    protected static final int PRIMARY_REPLICATE_PORT = 1237;
    protected static final int REPLICA_RPC_PORT = 4321;
    protected static final int REPLICA_PUSH_PORT = 4322;
    protected static final int REPLICA_FETCH_PORT = 4323;
    protected static final int REPLICA_REPLICATE_PORT = 4324;
    protected static final PartitionLocation primaryLocation =
        new PartitionLocation(
            0,
            1,
            "localhost",
            PRIMARY_RPC_PORT,
            PRIMARY_PUSH_PORT,
            PRIMARY_FETCH_PORT,
            PRIMARY_REPLICATE_PORT,
            PartitionLocation.Mode.PRIMARY);
    protected static final PartitionLocation replicaLocation =
        new PartitionLocation(
            0,
            1,
            "localhost",
            REPLICA_RPC_PORT,
            REPLICA_PUSH_PORT,
            REPLICA_FETCH_PORT,
            REPLICA_REPLICATE_PORT,
            PartitionLocation.Mode.REPLICA);
  
    protected final int BATCH_HEADER_SIZE = 4 * 4;
    protected ChannelFuture mockedFuture = mock(ChannelFuture.class);
  
    protected CelebornConf setupEnv(CompressionCodec codec) throws IOException, InterruptedException {
      CelebornConf conf = new CelebornConf();
      conf.set(CelebornConf.SHUFFLE_COMPRESSION_CODEC().key(), codec.name());
      conf.set(CelebornConf.CLIENT_PUSH_RETRY_THREADS().key(), "1");
      conf.set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "1K");
      shuffleClient =
          new ShuffleClientImpl(TEST_APPLICATION_ID, conf, new UserIdentifier("mock", "mock"));
      primaryLocation.setPeer(replicaLocation);
  
      when(endpointRef.askSync(
              ControlMessages.RegisterShuffle$.MODULE$.apply(TEST_SHUFFLE_ID, 1, 1),
              ClassTag$.MODULE$.apply(PbRegisterShuffleResponse.class)))
          .thenAnswer(
              t ->
                  ControlMessages.RegisterShuffleResponse$.MODULE$.apply(
                      StatusCode.SUCCESS, new PartitionLocation[] {primaryLocation}));
  
      shuffleClient.setupLifecycleManagerRef(endpointRef);
      when(clientFactory.createClient(
              primaryLocation.getHost(), primaryLocation.getPushPort(), TEST_REDUCRE_ID))
          .thenAnswer(t -> client);
  
      shuffleClient.dataClientFactory = clientFactory;
      return conf;
    }
  }
  