package com.atlassian.sal.core.executor;
   
   import com.atlassian.sal.api.executor.ThreadLocalContextManager;
   import com.atlassian.sal.api.executor.ThreadLocalDelegateExecutorFactory;
   import org.junit.Before;
   import org.junit.Test;
   import org.junit.runner.RunWith;
   import org.mockito.Mock;
   import org.mockito.invocation.InvocationOnMock;
  import org.mockito.junit.MockitoJUnitRunner;
  import org.mockito.stubbing.Answer;
  
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.Executor;
  import java.util.concurrent.Executors;
  import java.util.concurrent.TimeUnit;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  import static org.junit.Assert.assertNotNull;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.Mockito.when;
  
  @RunWith(MockitoJUnitRunner.class)
  public final class TestThreadLocalDelegateExecutor {
  
      @Mock
      private ThreadLocalDelegateExecutorFactory delegateExecutorFactory;
  
      @Before
      public void setup() {
          when(delegateExecutorFactory.createRunnable(any(Runnable.class))).thenAnswer(new Answer<Runnable>() {
              @Override
              public Runnable answer(InvocationOnMock invocation) throws Throwable {
                  return invocation.getArgument(0);
              }
          });
      }
  
      @Test
      public void testRun() throws InterruptedException {
          final ThreadLocalContextManager<Object> manager = new StubThreadLocalContextManager();
          Runnable delegate = new Runnable() {
              public void run() {
                  assertNotNull(manager.getThreadLocalContext());
              }
          };
  
          manager.setThreadLocalContext(new Object());
          final Executor executor = new ThreadLocalDelegateExecutor(Executors.newSingleThreadExecutor(), delegateExecutorFactory);
  
          final CountDownLatch latch = new CountDownLatch(1);
  
          executor.execute(new CountingDownRunnable(latch, delegate));
  
          latch.await(2, TimeUnit.SECONDS);
          assertNotNull(manager.getThreadLocalContext());
      }
  
      private static final class CountingDownRunnable implements Runnable {
          private final CountDownLatch latch;
          private final Runnable delegate;
  
          private CountingDownRunnable(CountDownLatch latch, Runnable delegate) {
              this.latch = checkNotNull(latch);
              this.delegate = checkNotNull(delegate);
          }
  
          @Override
          public void run() {
              try {
                  delegate.run();
              } finally {
                  latch.countDown();
              }
          }
      }
  }