ThriftStatsExporterInterceptorTest.java

/**
 * 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.thrift.aop;

import com.google.common.collect.ImmutableSet;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.matcher.Matchers;

import org.aopalliance.intercept.MethodInvocation;
import org.apache.aurora.common.stats.Stats;
import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.gen.GetJobsResult;
import org.apache.aurora.gen.JobConfiguration;
import org.apache.aurora.gen.Response;
import org.apache.aurora.gen.Result;
import org.apache.aurora.scheduler.thrift.auth.DecoratedThrift;
import org.junit.Before;
import org.junit.Test;

import static org.apache.aurora.gen.ResponseCode.OK;
import static org.apache.aurora.scheduler.thrift.Responses.error;
import static org.apache.aurora.scheduler.thrift.Responses.ok;
import static org.easymock.EasyMock.expect;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.fail;

public class ThriftStatsExporterInterceptorTest extends EasyMockTest {

  private static final String ROLE = "bob";

  private AnnotatedAuroraAdmin realThrift;
  private AnnotatedAuroraAdmin decoratedThrift;
  private ThriftStatsExporterInterceptor statsInterceptor;

  @Before
  public void setUp() {
    statsInterceptor = new ThriftStatsExporterInterceptor();
    realThrift = createMock(AnnotatedAuroraAdmin.class);
    Injector injector = Guice.createInjector(new AbstractModule() {
      @Override
      protected void configure() {
        MockDecoratedThrift.bindForwardedMock(binder(), realThrift);
        AopModule.bindThriftDecorator(
            binder(),
            Matchers.annotatedWith(DecoratedThrift.class),
            statsInterceptor);
      }
    });
    decoratedThrift = injector.getInstance(AnnotatedAuroraAdmin.class);
    addTearDown(Stats::flush);
  }

  @Test
  public void testIncrementStat() throws Exception {
    Response response = new Response().setResponseCode(OK)
        .setResult(Result.getJobsResult(new GetJobsResult()
        .setConfigs(ImmutableSet.of())));

    expect(realThrift.getJobs(ROLE)).andReturn(response);
    control.replay();

    assertSame(response, decoratedThrift.getJobs(ROLE));
    String statName = timingStatName("getJobs");
    assertEquals(1L, Stats.getVariable(statName + "_events").read());
    assertNotNull(Stats.getVariable(statName + "_events_per_sec"));
    assertNotNull(Stats.getVariable(statName + "_nanos_per_event"));
    assertNotNull(Stats.getVariable(statName + "_nanos_total"));
    assertNotNull(Stats.getVariable(statName + "_nanos_total_per_sec"));
  }

  @Test
  public void testMeasuredMethod() throws Throwable {
    MethodInvocation invocation = createMock(MethodInvocation.class);

    expect(invocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("measuredMethod"));
    expect(invocation.proceed()).andReturn(ok().setResult(Result.getJobsResult(
        new GetJobsResult(ImmutableSet.of(new JobConfiguration())))));

    control.replay();

    new ThriftStatsExporterInterceptor().invoke(invocation);
    assertEquals(1L, Stats.getVariable(workloadStatName("measuredMethod")).read());
  }

  @Test
  public void testUnmeasuredMethod() throws Throwable {
    MethodInvocation invocation = createMock(MethodInvocation.class);

    expect(invocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("unmeasuredMethod"));
    expect(invocation.proceed()).andReturn(ok().setResult(Result.getJobsResult(
        new GetJobsResult(ImmutableSet.of(new JobConfiguration())))));

    control.replay();

    new ThriftStatsExporterInterceptor().invoke(invocation);
    assertNull(Stats.getVariable(workloadStatName("unmeasuredMethod")));
  }

  @Test
  public void testExceptionalMeasuredMethod() throws Throwable {
    MethodInvocation invocation = createMock(MethodInvocation.class);

    expect(invocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("measuredMethod"));
    expect(invocation.proceed()).andThrow(new Exception());

    control.replay();

    try {
      new ThriftStatsExporterInterceptor().invoke(invocation);
      fail("Should not be reached");
    } catch (Exception e) {
      assertNull(Stats.getVariable(workloadStatName("measuredMethod")));
    }
  }

  @Test
  public void testMeasuredMethodWithErrorResponse() throws Throwable {
    MethodInvocation invocation = createMock(MethodInvocation.class);

    expect(invocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("measuredMethod"));
    expect(invocation.proceed()).andReturn(error("ERROR"));

    control.replay();

    new ThriftStatsExporterInterceptor().invoke(invocation);
    assertNull(Stats.getVariable(workloadStatName("measuredMethod")));
  }

  private static class InterceptedClass {
    @ThriftWorkload
    public Response measuredMethod() {
      throw new UnsupportedOperationException("Should not be called.");
    }

    public Response unmeasuredMethod() {
      throw new UnsupportedOperationException("Should not be called.");
    }
  }

  private static String timingStatName(String methodName) {
    return statName(ThriftStatsExporterInterceptor.TIMING_STATS_NAME_TEMPLATE, methodName);
  }

  private static String workloadStatName(String methodName) {
    return statName(ThriftStatsExporterInterceptor.WORKLOAD_STATS_NAME_TEMPLATE, methodName);
  }

  private static String statName(String template, String methodName) {
    return String.format(template, methodName);
  }

}