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.arrow.flight.client;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  
  import com.google.common.collect.ImmutableMap;
  import java.util.HashMap;
  import java.util.Map;
  import org.apache.arrow.flight.Action;
  import org.apache.arrow.flight.CallHeaders;
  import org.apache.arrow.flight.CallInfo;
  import org.apache.arrow.flight.CallStatus;
  import org.apache.arrow.flight.Criteria;
  import org.apache.arrow.flight.FlightCallHeaders;
  import org.apache.arrow.flight.FlightClient;
  import org.apache.arrow.flight.FlightClient.ClientStreamListener;
  import org.apache.arrow.flight.FlightDescriptor;
  import org.apache.arrow.flight.FlightMethod;
  import org.apache.arrow.flight.FlightServer;
  import org.apache.arrow.flight.FlightServerMiddleware;
  import org.apache.arrow.flight.FlightStream;
  import org.apache.arrow.flight.FlightTestUtil;
  import org.apache.arrow.flight.HeaderCallOption;
  import org.apache.arrow.flight.Location;
  import org.apache.arrow.flight.NoOpFlightProducer;
  import org.apache.arrow.flight.RequestContext;
  import org.apache.arrow.flight.SyncPutListener;
  import org.apache.arrow.flight.Ticket;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.util.AutoCloseables;
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  /** Tests to ensure custom headers are passed along to the server for each command. */
  public class CustomHeaderTest {
    FlightServer server;
    FlightClient client;
    BufferAllocator allocator;
    TestCustomHeaderMiddleware.Factory headersMiddleware;
    HeaderCallOption headers;
    Map<String, String> testHeaders =
        ImmutableMap.of(
            "foo", "bar",
            "bar", "foo",
            "answer", "42");
  
    @BeforeEach
    public void setUp() throws Exception {
      allocator = new RootAllocator(Integer.MAX_VALUE);
      headersMiddleware = new TestCustomHeaderMiddleware.Factory();
      FlightCallHeaders callHeaders = new FlightCallHeaders();
      for (Map.Entry<String, String> entry : testHeaders.entrySet()) {
        callHeaders.insert(entry.getKey(), entry.getValue());
      }
      headers = new HeaderCallOption(callHeaders);
      server =
          FlightServer.builder(
                  allocator,
                  Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, /*port*/ 0),
                  new NoOpFlightProducer())
              .middleware(FlightServerMiddleware.Key.of("customHeader"), headersMiddleware)
              .build();
      server.start();
      client = FlightClient.builder(allocator, server.getLocation()).build();
    }
  
    @AfterEach
    public void tearDown() throws Exception {
      allocator.getChildAllocators().forEach(BufferAllocator::close);
      AutoCloseables.close(allocator, server, client);
    }
  
    @Test
    public void testHandshake() {
      try {
        client.handshake(headers);
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.HANDSHAKE);
    }
  
    @Test
    public void testGetSchema() {
      try {
        client.getSchema(FlightDescriptor.command(new byte[0]), headers);
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.GET_SCHEMA);
    }
  
    @Test
    public void testGetFlightInfo() {
      try {
        client.getInfo(FlightDescriptor.command(new byte[0]), headers);
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.GET_FLIGHT_INFO);
    }
  
    @Test
    public void testListActions() {
      try {
        client.listActions(headers).iterator().next();
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.LIST_ACTIONS);
    }
  
    @Test
    public void testListFlights() {
      try {
        client.listFlights(new Criteria(new byte[] {1}), headers).iterator().next();
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.LIST_FLIGHTS);
    }
  
    @Test
    public void testDoAction() {
      try {
        client.doAction(new Action("test"), headers).next();
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.DO_ACTION);
    }
  
    @Test
    public void testStartPut() {
      try {
        final ClientStreamListener listener =
            client.startPut(FlightDescriptor.command(new byte[0]), new SyncPutListener(), headers);
        listener.getResult();
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.DO_PUT);
    }
  
    @Test
    public void testGetStream() {
      try (final FlightStream stream = client.getStream(new Ticket(new byte[0]), headers)) {
        stream.next();
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.DO_GET);
    }
  
    @Test
    public void testDoExchange() {
      try (final FlightClient.ExchangeReaderWriter stream =
          client.doExchange(FlightDescriptor.command(new byte[0]), headers)) {
        stream.getReader().next();
      } catch (Exception ignored) {
      }
  
      assertHeadersMatch(FlightMethod.DO_EXCHANGE);
    }
  
    private void assertHeadersMatch(FlightMethod method) {
      for (Map.Entry<String, String> entry : testHeaders.entrySet()) {
        assertEquals(entry.getValue(), headersMiddleware.getCustomHeader(method, entry.getKey()));
      }
    }
  
    /** A middleware used to test if customHeaders are being sent to the server properly. */
    static class TestCustomHeaderMiddleware implements FlightServerMiddleware {
  
      public TestCustomHeaderMiddleware() {}
  
      @Override
      public void onBeforeSendingHeaders(CallHeaders callHeaders) {}
  
      @Override
      public void onCallCompleted(CallStatus callStatus) {}
  
      @Override
      public void onCallErrored(Throwable throwable) {}
  
      /**
       * A factory for the middleware that keeps track of the received headers and provides a way to
       * check those values for a given Flight Method.
       */
      static class Factory implements FlightServerMiddleware.Factory<TestCustomHeaderMiddleware> {
        private final Map<FlightMethod, CallHeaders> receivedCallHeaders = new HashMap<>();
  
        @Override
        public TestCustomHeaderMiddleware onCallStarted(
            CallInfo callInfo, CallHeaders callHeaders, RequestContext requestContext) {
  
          receivedCallHeaders.put(callInfo.method(), callHeaders);
          return new TestCustomHeaderMiddleware();
        }
  
        public String getCustomHeader(FlightMethod method, String key) {
          CallHeaders headers = receivedCallHeaders.get(method);
          if (headers == null) {
            return null;
          }
          return headers.get(key);
        }
      }
    }
  }
  