package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalContextManager;
   import com.atlassian.sal.api.executor.ThreadLocalDelegateExecutorFactory;
   import com.google.common.collect.ImmutableSet;
   import org.junit.Test;
   
   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;
  
  import static org.junit.Assert.assertEquals;
  
  /**
   * 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 {
      private static final Object OUTER_THREAD_LOCAL_CONTEXT = new Object();
  
      private final ThreadLocalContextManager<Object> manager = threadLocalContextManager();
      private final ThreadLocalDelegateExecutorFactory executorFactory = new DefaultThreadLocalDelegateExecutorFactory<>(manager);
  
      @Test
      public void testSubmitCallable() throws ExecutionException, InterruptedException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          service.submit(assertingCallable()).get();
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      public void testSubmitRunnableWithResult() throws ExecutionException, InterruptedException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          service.submit(assertingRunnable(), null).get();
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      public void testSubmitRunnable() throws ExecutionException, InterruptedException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          service.submit(assertingRunnable()).get();
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      public void testInvokeAll() throws ExecutionException, InterruptedException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          getAllFutures(service.invokeAll(assertingCallables()));
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      public void testInvokeAllWithTimeout() throws ExecutionException, InterruptedException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          getAllFutures(service.invokeAll(assertingCallables(), 60, TimeUnit.SECONDS));
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      public void testInvokeAny() throws ExecutionException, InterruptedException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          service.invokeAny(assertingCallables());
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      public void testInvokeAnyWithTimeout() throws ExecutionException, InterruptedException, TimeoutException {
          final ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          service.invokeAny(assertingCallables(), 60, TimeUnit.SECONDS);
          assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
      }
  
      @Test
      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 ThreadLocalDelegateExecutorService service = threadLocalDelegateExecutorService();
          final AtomicReference<AssertionError> assertionErrorReference = new AtomicReference<>();
  
          service.execute(assertingRunnableWithErrorReference(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 threadLocalDelegateExecutorService() {
         return new ThreadLocalDelegateExecutorService(Executors.newSingleThreadExecutor(), executorFactory);
     }
 
     private Runnable assertingRunnable() {
         return new Runnable() {
             @Override
             public void run() {
                 assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
             }
         };
     }
 
     private Runnable assertingRunnableWithErrorReference(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() {
         return new Callable<Object>() {
             @Override
             public Object call() throws Exception {
                 assertEquals(OUTER_THREAD_LOCAL_CONTEXT, manager.getThreadLocalContext());
                 return null;
             }
         };
     }
 
     private Collection<Callable<Object>> assertingCallables() {
         return ImmutableSet.of(assertingCallable(), assertingCallable(), assertingCallable());
     }
 
     private void getAllFutures(Collection<Future<Object>> futures) throws ExecutionException, InterruptedException {
         for (Future<Object> future : futures) {
             future.get();
         }
     }
 }