package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalContextManager;
   import com.google.common.collect.ImmutableSet;
   import junit.framework.TestCase;
   
   import java.util.Collection;
   import java.util.concurrent.Callable;
   import java.util.concurrent.ExecutionException;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.concurrent.atomic.AtomicReference;
  
  /**
   * Test that ThreadLocalDelegateExecutorService correctly wraps callables and runnables passed to it before passing them
   * off to its delegate, by checking the thread-local context in the call or run method. Also, test that the thread-local
   * context doesn't get set to something else on return.
   */
  public class TestThreadLocalDelegateExecutorService extends TestCase {
      private static final Object OUTER_THREAD_LOCAL_CONTEXT = new Object();
  
      public void testSubmitCallable() throws ExecutionException, InterruptedException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          service.submit(assertingCallable(manager)).get();
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testSubmitRunnableWithResult() throws ExecutionException, InterruptedException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          service.submit(assertingRunnable(manager), null).get();
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testSubmitRunnable() throws ExecutionException, InterruptedException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService(manager);
          service.submit(assertingRunnable(manager)).get();
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testInvokeAll() throws ExecutionException, InterruptedException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          getAllFutures(service.invokeAll(assertingCallables(manager)));
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testInvokeAllWithTimeout() throws ExecutionException, InterruptedException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          getAllFutures(service.invokeAll(assertingCallables(manager), 60, TimeUnit.SECONDS));
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testInvokeAny() throws ExecutionException, InterruptedException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          service.invokeAny(assertingCallables(manager));
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testInvokeAnyWithTimeout() throws ExecutionException, InterruptedException, TimeoutException {
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          service.invokeAny(assertingCallables(manager), 60, TimeUnit.SECONDS);
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      public void testExecute() throws Exception {
          // We have to jump through a few hoops here, since the AssertionError gets thrown on a different thread, and
          // there is no Future to get.
  
          final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
          final ThreadLocalDelegateExecutorService<Object> service = threadLocalDelegateExecutorService(manager);
          final AtomicReference<AssertionError> assertionErrorReference = new AtomicReference<AssertionError>();
  
          service.execute(assertingRunnableWithErrorReference(manager, assertionErrorReference));
          service.shutdown();
          service.awaitTermination(60, TimeUnit.SECONDS);
  
          final AssertionError assertionError = assertionErrorReference.get();
          if (assertionError != null)
              throw assertionError;
  
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      private ThreadLocalContextManager<Object> threadLocalContextManager() {
          ThreadLocalContextManager<Object> manager = new StubThreadLocalContextManager();
          manager.setThreadLocalContext(OUTER_THREAD_LOCAL_CONTEXT);
          return manager;
      }
  
      private ThreadLocalDelegateExecutorService<Object> threadLocalDelegateExecutorService(ThreadLocalContextManager<Object> manager) {
          return new ThreadLocalDelegateExecutorService<Object>(manager, Executors.newSingleThreadExecutor());
     }
 
     private Runnable assertingRunnable(final ThreadLocalContextManager manager) {
         return new Runnable() {
             @Override
             public void run() {
                 assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
             }
         };
     }
 
     private Runnable assertingRunnableWithErrorReference(final ThreadLocalContextManager manager, final AtomicReference<AssertionError> assertionErrorReference) {
         return new Runnable() {
             @Override
             public void run() {
                 try {
                     assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
                 } catch (AssertionError e) {
                     assertionErrorReference.set(e);
                 }
             }
         };
     }
 
     private Callable<Object> assertingCallable(final ThreadLocalContextManager manager) {
         return new Callable<Object>() {
             @Override
             public Object call() throws Exception {
                 assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
                 return null;
             }
         };
     }
 
     private Collection<Callable<Object>> assertingCallables(final ThreadLocalContextManager manager) {
         return ImmutableSet.of(assertingCallable(manager), assertingCallable(manager), assertingCallable(manager));
     }
 
     private void getAllFutures(Collection<Future<Object>> futures) throws ExecutionException, InterruptedException {
         for (Future<Object> future : futures) {
             future.get();
         }
     }
 }