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 String OUTER_THREAD_LOCAL_CONTEXT = "OUTER_THREAD_LOCAL_CONTEXT";
  
      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 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);
  
          if (assertionErrorReference.get() != null)
              throw assertionErrorReference.get();
  
          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();
         }
     }
 }