content: /**
   * Licensed 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.aurora.scheduler.http;
  
  import java.util.Arrays;
  import java.util.Optional;
  import java.util.concurrent.atomic.AtomicReference;
  
  import javax.servlet.http.HttpServletRequest;
  
  import com.google.common.base.Function;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Iterables;
  import com.google.common.net.HostAndPort;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.scheduler.app.ServiceGroupMonitor;
  import org.apache.aurora.scheduler.app.ServiceGroupMonitor.MonitorException;
  import org.apache.aurora.scheduler.discovery.ServiceInstance;
  import org.apache.aurora.scheduler.discovery.ServiceInstance.Endpoint;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.apache.aurora.scheduler.http.LeaderRedirect.LeaderStatus;
  import static org.easymock.EasyMock.expect;
  import static org.easymock.EasyMock.expectLastCall;
  import static org.junit.Assert.assertEquals;
  
  public class LeaderRedirectTest extends EasyMockTest {
  
    private static final int HTTP_PORT = 500;
  
    private static final Function<HostAndPort, ServiceInstance> CREATE_INSTANCE =
        endpoint -> new ServiceInstance(
            new Endpoint(endpoint.getHost(), endpoint.getPort()),
            ImmutableMap.of());
  
    private AtomicReference<ImmutableSet<ServiceInstance>> schedulers;
    private ServiceGroupMonitor serviceGroupMonitor;
    private LeaderRedirect leaderRedirector;
  
    @Before
    public void setUp() throws MonitorException {
      schedulers = new AtomicReference<>(ImmutableSet.of());
      serviceGroupMonitor = createMock(ServiceGroupMonitor.class);
  
      HttpService http = createMock(HttpService.class);
      expect(http.getAddress()).andStubReturn(HostAndPort.fromParts("localhost", HTTP_PORT));
  
      leaderRedirector = new LeaderRedirect(http, serviceGroupMonitor);
    }
  
    private void replayAndMonitor(int expectedGetCalls) throws Exception {
      serviceGroupMonitor.start();
      expectLastCall();
  
      expect(serviceGroupMonitor.get()).andAnswer(() -> schedulers.get()).times(expectedGetCalls);
  
      control.replay();
      leaderRedirector.monitor();
    }
  
    @Test
    public void testLeader() throws Exception {
      replayAndMonitor(3);
      publishSchedulers(localPort(HTTP_PORT));
  
      assertEquals(Optional.empty(), leaderRedirector.getRedirect());
  
      // NB: LEADING takes 2 tests of the server group membership to calculate; thus we expect 3
      // server group get calls, 1 for the getRedirect() above and 2 here.
      assertEquals(LeaderStatus.LEADING, leaderRedirector.getLeaderStatus());
    }
  
    @Test
    public void testNotLeader() throws Exception {
      replayAndMonitor(3);
  
      HostAndPort remote = HostAndPort.fromParts("foobar", HTTP_PORT);
      publishSchedulers(remote);
  
      assertEquals(Optional.of(remote), leaderRedirector.getRedirect());
  
      // NB: NOT_LEADING takes 2 tests of the server group membership to calculate; thus we expect 3
      // server group get calls, 1 for the getRedirect() above and 2 here.
      assertEquals(LeaderStatus.NOT_LEADING, leaderRedirector.getLeaderStatus());
    }
  
    @Test
    public void testLeaderOnSameHost() throws Exception {
      replayAndMonitor(3);
  
      HostAndPort local = localPort(555);
      publishSchedulers(local);
  
      assertEquals(Optional.of(local), leaderRedirector.getRedirect());
  
      // NB: NOT_LEADING takes 2 tests of the server group membership to calculate; thus we expect 3
      // server group get calls, 1 for the getRedirect() above and 2 here.
      assertEquals(LeaderStatus.NOT_LEADING, leaderRedirector.getLeaderStatus());
    }
  
    @Test
    public void testNoLeaders() throws Exception {
      replayAndMonitor(2);
  
      assertEquals(Optional.empty(), leaderRedirector.getRedirect());
      assertEquals(LeaderStatus.NO_LEADER, leaderRedirector.getLeaderStatus());
    }
  
    @Test
    public void testMultipleLeaders() throws Exception {
      replayAndMonitor(2);
  
      publishSchedulers(HostAndPort.fromParts("foobar", 500), HostAndPort.fromParts("baz", 800));
  
      assertEquals(Optional.empty(), leaderRedirector.getRedirect());
      assertEquals(LeaderStatus.NO_LEADER, leaderRedirector.getLeaderStatus());
    }
  
    private HttpServletRequest mockRequest(String attributeValue, String queryString) {
      HttpServletRequest mockRequest = createMock(HttpServletRequest.class);
      expect(mockRequest.getScheme()).andReturn("http");
      expect(mockRequest.getAttribute(JettyServerModule.ORIGINAL_PATH_ATTRIBUTE_NAME))
          .andReturn(attributeValue);
      expect(mockRequest.getRequestURI()).andReturn("/some/path");
      expect(mockRequest.getQueryString()).andReturn(queryString);
  
      return mockRequest;
    }
  
    @Test
    public void testRedirectTargetNoAttribute() throws Exception {
      HttpServletRequest mockRequest = mockRequest(null, null);
  
      replayAndMonitor(1);
  
      HostAndPort remote = HostAndPort.fromParts("foobar", HTTP_PORT);
      publishSchedulers(remote);
  
      assertEquals(
          Optional.of("http://foobar:500/some/path"),
          leaderRedirector.getRedirectTarget(mockRequest));
    }
  
    @Test
    public void testRedirectTargetWithAttribute() throws Exception {
      HttpServletRequest mockRequest = mockRequest("/the/original/path", null);
  
      replayAndMonitor(1);
  
      HostAndPort remote = HostAndPort.fromParts("foobar", HTTP_PORT);
      publishSchedulers(remote);
  
      assertEquals(
          Optional.of("http://foobar:500/the/original/path"),
          leaderRedirector.getRedirectTarget(mockRequest));
    }
  
    @Test
    public void testRedirectTargetQueryString() throws Exception {
      HttpServletRequest mockRequest = mockRequest(null, "bar=baz");
  
      replayAndMonitor(1);
  
      HostAndPort remote = HostAndPort.fromParts("foobar", HTTP_PORT);
      publishSchedulers(remote);
  
      assertEquals(
          Optional.of("http://foobar:500/some/path?bar=baz"),
          leaderRedirector.getRedirectTarget(mockRequest));
    }
  
    private void publishSchedulers(HostAndPort... schedulerHttpEndpoints) {
      publishSchedulers(ImmutableSet.copyOf(Iterables.transform(Arrays.asList(schedulerHttpEndpoints),
          CREATE_INSTANCE)));
    }
  
    private void publishSchedulers(ImmutableSet<ServiceInstance> instances) {
      schedulers.set(instances);
    }
  
    private static HostAndPort localPort(int port) {
      return HostAndPort.fromParts("localhost", port);
    }
  }
  