LoggingInterceptorTest.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 java.util.concurrent.atomic.AtomicBoolean;

import javax.annotation.Nullable;

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.JobConfiguration;
import org.apache.aurora.gen.Response;
import org.apache.aurora.gen.ResponseCode;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.scheduler.storage.Storage.TransientStorageException;
import org.junit.Before;
import org.junit.Test;

import static org.easymock.EasyMock.expect;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;

public class LoggingInterceptorTest extends EasyMockTest {
  private MethodInvocation methodInvocation;

  private LoggingInterceptor loggingInterceptor;

  // java.lang.reflect.Method is final so we need to create an actual class to get instances of
  // them.
  // TODO(ksweeney): Investigate using PowerMock to do this instead.
  private static class InterceptedClass {
    Response respond() {
      throw new UnsupportedOperationException();
    }

    Response respond(Object arg1, @Nullable Object arg2) {
      throw new UnsupportedOperationException();
    }

    Response respond(JobConfiguration jobConfiguration) {
      throw new UnsupportedOperationException();
    }
  }

  @Before
  public void setUp() throws Exception {
    loggingInterceptor = new LoggingInterceptor();

    methodInvocation = createMock(MethodInvocation.class);
  }

  @Test
  public void testInvokeTransientStorageException() throws Throwable {
    expect(methodInvocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("respond"));
    expect(methodInvocation.getArguments()).andReturn(new Object[]{});
    expect(methodInvocation.proceed()).andThrow(new TransientStorageException("try again"));

    control.replay();

    assertNull(Stats.getVariable("scheduler_thrift_response_ERROR_TRANSIENT"));
    assertEquals(
        ResponseCode.ERROR_TRANSIENT,
        ((Response) loggingInterceptor.invoke(methodInvocation)).getResponseCode());
    assertEquals(1L, Stats.getVariable("scheduler_thrift_response_ERROR_TRANSIENT").read());
  }

  @Test
  public void testInvokeRuntimeException() throws Throwable {
    expect(methodInvocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("respond"));
    expect(methodInvocation.getArguments()).andReturn(new Object[] {});
    expect(methodInvocation.proceed()).andThrow(new RuntimeException());

    control.replay();

    assertNull(Stats.getVariable("scheduler_thrift_response_ERROR"));
    assertEquals(
        ResponseCode.ERROR,
        ((Response) loggingInterceptor.invoke(methodInvocation)).getResponseCode());
    assertEquals(1L, Stats.getVariable("scheduler_thrift_response_ERROR").read());
  }

  @Test
  public void testInvokePrintsArgs() throws Throwable {
    Response response = new Response();
    response.setResponseCode(ResponseCode.OK);

    expect(methodInvocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("respond", Object.class, Object.class));

    // EasyMock cannot mock toString.
    final AtomicBoolean calledAtLeastOnce = new AtomicBoolean(false);
    Object arg1 = new Object() {
      @Override
      public String toString() {
        calledAtLeastOnce.set(true);
        return "arg1";
      }
    };
    expect(methodInvocation.getArguments()).andReturn(new Object[] {arg1, null});
    expect(methodInvocation.proceed()).andReturn(response);

    control.replay();

    assertNull(Stats.getVariable("scheduler_thrift_response_OK"));
    assertSame(response, loggingInterceptor.invoke(methodInvocation));
    assertTrue(calledAtLeastOnce.get());
    assertEquals(1L, Stats.getVariable("scheduler_thrift_response_OK").read());
  }

  @Test
  public void testInvokePrintsBlankedJobConfiguration() throws Throwable {
    Response response = new Response();
    response.setResponseCode(ResponseCode.INVALID_REQUEST);

    expect(methodInvocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("respond", JobConfiguration.class));
    expect(methodInvocation.getArguments())
        .andReturn(new Object[] {new JobConfiguration().setTaskConfig(new TaskConfig())});
    expect(methodInvocation.proceed()).andReturn(response);

    control.replay();

    assertNull(Stats.getVariable("scheduler_thrift_response_INVALID_REQUEST"));
    assertSame(response, loggingInterceptor.invoke(methodInvocation));
    assertEquals(1L, Stats.getVariable("scheduler_thrift_response_INVALID_REQUEST").read());
  }

  @Test
  public void testInvokePrintsJobConfiguration() throws Throwable {
    Response response = new Response();
    response.setResponseCode(ResponseCode.WARNING);

    expect(methodInvocation.getMethod())
        .andReturn(InterceptedClass.class.getDeclaredMethod("respond", JobConfiguration.class));
    expect(methodInvocation.getArguments()).andReturn(new Object[] {new JobConfiguration()});
    expect(methodInvocation.proceed()).andReturn(response);

    control.replay();

    assertNull(Stats.getVariable("scheduler_thrift_response_WARNING"));
    assertSame(response, loggingInterceptor.invoke(methodInvocation));
    assertEquals(1L, Stats.getVariable("scheduler_thrift_response_WARNING").read());
  }
}