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.kafka.clients.admin.internals;
  
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.errors.DisconnectException;
  import org.apache.kafka.common.errors.UnknownServerException;
  import org.apache.kafka.common.internals.KafkaFutureImpl;
  import org.apache.kafka.common.message.MetadataRequestData;
  import org.apache.kafka.common.message.MetadataResponseData;
  import org.apache.kafka.common.protocol.ApiKeys;
  import org.apache.kafka.common.requests.AbstractRequest;
  import org.apache.kafka.common.requests.AbstractResponse;
  import org.apache.kafka.common.requests.MetadataRequest;
  import org.apache.kafka.common.requests.MetadataResponse;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.MockTime;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.Test;
  
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.OptionalInt;
  import java.util.Set;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotEquals;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  public class AllBrokersStrategyIntegrationTest {
      private static final long TIMEOUT_MS = 5000;
      private static final long RETRY_BACKOFF_MS = 100;
  
      private final LogContext logContext = new LogContext();
      private final MockTime time = new MockTime();
  
      private AdminApiDriver<AllBrokersStrategy.BrokerKey, Integer> buildDriver(
          AllBrokersStrategy.AllBrokersFuture<Integer> result
      ) {
          return new AdminApiDriver<>(
              new MockApiHandler(),
              result,
              time.milliseconds() + TIMEOUT_MS,
              RETRY_BACKOFF_MS,
              RETRY_BACKOFF_MS,
              logContext
          );
      }
  
      @Test
      public void testFatalLookupError() {
          AllBrokersStrategy.AllBrokersFuture<Integer> result = new AllBrokersStrategy.AllBrokersFuture<>();
          AdminApiDriver<AllBrokersStrategy.BrokerKey, Integer> driver = buildDriver(result);
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> requestSpecs = driver.poll();
          assertEquals(1, requestSpecs.size());
  
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> spec = requestSpecs.get(0);
          assertEquals(AllBrokersStrategy.LOOKUP_KEYS, spec.keys);
  
          driver.onFailure(time.milliseconds(), spec, new UnknownServerException());
          assertTrue(result.all().isDone());
          TestUtils.assertFutureThrows(UnknownServerException.class, result.all());
          assertEquals(Collections.emptyList(), driver.poll());
      }
  
      @Test
      public void testRetryLookupAfterDisconnect() {
          AllBrokersStrategy.AllBrokersFuture<Integer> result = new AllBrokersStrategy.AllBrokersFuture<>();
          AdminApiDriver<AllBrokersStrategy.BrokerKey, Integer> driver = buildDriver(result);
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> requestSpecs = driver.poll();
          assertEquals(1, requestSpecs.size());
  
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> spec = requestSpecs.get(0);
          assertEquals(AllBrokersStrategy.LOOKUP_KEYS, spec.keys);
  
          driver.onFailure(time.milliseconds(), spec, new DisconnectException());
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> retrySpecs = driver.poll();
          assertEquals(1, retrySpecs.size());
  
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> retrySpec = retrySpecs.get(0);
          assertEquals(AllBrokersStrategy.LOOKUP_KEYS, retrySpec.keys);
          assertEquals(time.milliseconds(), retrySpec.nextAllowedTryMs);
          assertEquals(Collections.emptyList(), driver.poll());
      }
  
      @Test
      public void testMultiBrokerCompletion() throws Exception {
          AllBrokersStrategy.AllBrokersFuture<Integer> result = new AllBrokersStrategy.AllBrokersFuture<>();
          AdminApiDriver<AllBrokersStrategy.BrokerKey, Integer> driver = buildDriver(result);
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> lookupSpecs = driver.poll();
          assertEquals(1, lookupSpecs.size());
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> lookupSpec = lookupSpecs.get(0);
  
          Set<Integer> brokerIds = Set.of(1, 2);
          driver.onResponse(time.milliseconds(), lookupSpec, responseWithBrokers(brokerIds), Node.noNode());
          assertTrue(result.all().isDone());
  
          Map<Integer, KafkaFutureImpl<Integer>> brokerFutures = result.all().get();
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> requestSpecs = driver.poll();
          assertEquals(2, requestSpecs.size());
  
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> requestSpec1 = requestSpecs.get(0);
          assertTrue(requestSpec1.scope.destinationBrokerId().isPresent());
          int brokerId1 = requestSpec1.scope.destinationBrokerId().getAsInt();
          assertTrue(brokerIds.contains(brokerId1));
  
          driver.onResponse(time.milliseconds(), requestSpec1, null, Node.noNode());
          KafkaFutureImpl<Integer> future1 = brokerFutures.get(brokerId1);
          assertTrue(future1.isDone());
  
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> requestSpec2 = requestSpecs.get(1);
          assertTrue(requestSpec2.scope.destinationBrokerId().isPresent());
          int brokerId2 = requestSpec2.scope.destinationBrokerId().getAsInt();
          assertNotEquals(brokerId1, brokerId2);
          assertTrue(brokerIds.contains(brokerId2));
  
          driver.onResponse(time.milliseconds(), requestSpec2, null, Node.noNode());
          KafkaFutureImpl<Integer> future2 = brokerFutures.get(brokerId2);
          assertTrue(future2.isDone());
          assertEquals(Collections.emptyList(), driver.poll());
      }
  
      @Test
      public void testRetryFulfillmentAfterDisconnect() throws Exception {
          AllBrokersStrategy.AllBrokersFuture<Integer> result = new AllBrokersStrategy.AllBrokersFuture<>();
          AdminApiDriver<AllBrokersStrategy.BrokerKey, Integer> driver = buildDriver(result);
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> lookupSpecs = driver.poll();
          assertEquals(1, lookupSpecs.size());
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> lookupSpec = lookupSpecs.get(0);
  
          int brokerId = 1;
          driver.onResponse(time.milliseconds(), lookupSpec, responseWithBrokers(Collections.singleton(brokerId)), Node.noNode());
          assertTrue(result.all().isDone());
  
          Map<Integer, KafkaFutureImpl<Integer>> brokerFutures = result.all().get();
          KafkaFutureImpl<Integer> future = brokerFutures.get(brokerId);
          assertFalse(future.isDone());
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> requestSpecs = driver.poll();
          assertEquals(1, requestSpecs.size());
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> requestSpec = requestSpecs.get(0);
  
          driver.onFailure(time.milliseconds(), requestSpec, new DisconnectException());
          assertFalse(future.isDone());
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> retrySpecs = driver.poll();
          assertEquals(1, retrySpecs.size());
  
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> retrySpec = retrySpecs.get(0);
          assertEquals(time.milliseconds() + RETRY_BACKOFF_MS, retrySpec.nextAllowedTryMs);
          assertEquals(OptionalInt.of(brokerId), retrySpec.scope.destinationBrokerId());
  
          driver.onResponse(time.milliseconds(), retrySpec, null, new Node(brokerId, "host", 1234));
          assertTrue(future.isDone());
          assertEquals(brokerId, future.get());
          assertEquals(Collections.emptyList(), driver.poll());
      }
  
      @Test
      public void testFatalFulfillmentError() throws Exception {
          AllBrokersStrategy.AllBrokersFuture<Integer> result = new AllBrokersStrategy.AllBrokersFuture<>();
          AdminApiDriver<AllBrokersStrategy.BrokerKey, Integer> driver = buildDriver(result);
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> lookupSpecs = driver.poll();
          assertEquals(1, lookupSpecs.size());
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> lookupSpec = lookupSpecs.get(0);
  
          int brokerId = 1;
          driver.onResponse(time.milliseconds(), lookupSpec, responseWithBrokers(Collections.singleton(brokerId)), Node.noNode());
          assertTrue(result.all().isDone());
  
          Map<Integer, KafkaFutureImpl<Integer>> brokerFutures = result.all().get();
          KafkaFutureImpl<Integer> future = brokerFutures.get(brokerId);
          assertFalse(future.isDone());
  
          List<AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey>> requestSpecs = driver.poll();
          assertEquals(1, requestSpecs.size());
          AdminApiDriver.RequestSpec<AllBrokersStrategy.BrokerKey> requestSpec = requestSpecs.get(0);
  
          driver.onFailure(time.milliseconds(), requestSpec, new UnknownServerException());
          assertTrue(future.isDone());
          TestUtils.assertFutureThrows(UnknownServerException.class, future);
          assertEquals(Collections.emptyList(), driver.poll());
      }
  
      private MetadataResponse responseWithBrokers(Set<Integer> brokerIds) {
          MetadataResponseData response = new MetadataResponseData();
          for (Integer brokerId : brokerIds) {
              response.brokers().add(new MetadataResponseData.MetadataResponseBroker()
                  .setNodeId(brokerId)
                  .setHost("host" + brokerId)
                  .setPort(9092)
              );
          }
          return new MetadataResponse(response, ApiKeys.METADATA.latestVersion());
      }
  
      private class MockApiHandler extends AdminApiHandler.Batched<AllBrokersStrategy.BrokerKey, Integer> {
          private final AllBrokersStrategy allBrokersStrategy = new AllBrokersStrategy(logContext);
  
          @Override
          public String apiName() {
              return "mock-api";
          }
  
          @Override
          public AbstractRequest.Builder<?> buildBatchedRequest(
              int brokerId,
              Set<AllBrokersStrategy.BrokerKey> keys
          ) {
              return new MetadataRequest.Builder(new MetadataRequestData());
          }
  
          @Override
          public ApiResult<AllBrokersStrategy.BrokerKey, Integer> handleResponse(
              Node broker,
              Set<AllBrokersStrategy.BrokerKey> keys,
              AbstractResponse response
          ) {
              return ApiResult.completed(keys.iterator().next(), broker.id());
          }
  
          @Override
          public AdminApiLookupStrategy<AllBrokersStrategy.BrokerKey> lookupStrategy() {
              return allBrokersStrategy;
          }
      }
  }
  