package org.springframework.cloud.sleuth.instrument;

import static org.assertj.core.api.BDDAssertions.then;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.junit.Ignore;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.mockito.runners.MockitoJUnitRunner;
import org.springframework.cloud.sleuth.Trace;
import org.springframework.cloud.sleuth.TraceManager;
import org.springframework.cloud.sleuth.trace.TraceContextHolder;

@RunWith(MockitoJUnitRunner.class)
public class TraceRunnableTest {

	ExecutorService executor = Executors.newSingleThreadExecutor();
	TraceManager traceManager = Mockito.mock(TraceManager.class);

	@Test
	@Ignore("Will fail because trace is not removed after runnable gets executed")
	public void should_remove_span_from_thread_local_after_finishing_work() throws Exception {
		// given
		TraceSettingRunnable traceSettingRunnable = runnableThatSetsTraceInCurrentThreadLocalWithInitialTrace();
		givenRunnableGetsSubmitted(traceSettingRunnable);
		then(traceSettingRunnable.trace).isNotNull();

		// when
		TraceKeepingRunnable traceKeepingRunnable = runnableThatRetrievesTraceFromThreadLocal();
		whenRunnableGetsSubmitted(traceKeepingRunnable);

		// then
		then(traceKeepingRunnable.trace).isNull();
	}

	private TraceSettingRunnable runnableThatSetsTraceInCurrentThreadLocalWithInitialTrace() {
		return new TraceSettingRunnable();
	}

	private TraceKeepingRunnable runnableThatRetrievesTraceFromThreadLocal() {
		return new TraceKeepingRunnable();
	}

	private void givenRunnableGetsSubmitted(Runnable runnable) throws Exception {
		whenRunnableGetsSubmitted(runnable);
	}

	private void whenRunnableGetsSubmitted(Runnable callable) throws Exception {
		executor.submit(new TraceRunnable(traceManager, callable)).get();
	}

	static class TraceKeepingRunnable implements Runnable {
		public Trace trace;

		@Override
		public void run() {
			trace = TraceContextHolder.getCurrentTrace();
		}
	}

	static class TraceSettingRunnable implements Runnable {
		public Trace trace;

		@Override
		public void run() {
			TraceContextHolder.setCurrentTrace(Mockito.mock(Trace.class));
			trace = TraceContextHolder.getCurrentTrace();
		}
	}
}
@@ -192,33 +192,8 @@
 s;%0A%0A
-import org.junit.Ignore;%0A
 impo
@@ -464,31 +464,271 @@
 uth.
-trace.TraceContextHolde
+sampler.AlwaysSampler;%0Aimport org.springframework.cloud.sleuth.trace.DefaultTraceManager;%0Aimport org.springframework.cloud.sleuth.trace.TraceContextHolder;%0Aimport org.springframework.context.ApplicationEventPublisher;%0Aimport org.springframework.util.JdkIdGenerato
 r;%0A%0A
@@ -894,129 +894,133 @@
 r = 
-Mockito.mock(TraceManager.class);%0A%0A%09@Test%0A%09@Ignore(%22Will fail because trace is not removed after runnable gets executed%22)
+new DefaultTraceManager(new AlwaysSampler(),%0A%09%09%09new JdkIdGenerator(), Mockito.mock(ApplicationEventPublisher.class));%0A%0A%09@Test
 %0A%09pu
@@ -1280,21 +1280,35 @@
 ble);%0A%09%09
-then(
+Trace firstTrace = 
 traceSet
@@ -1321,24 +1321,43 @@
 nnable.trace
+;%0A%09%09then(firstTrace
 ).isNotNull(
@@ -1525,21 +1525,36 @@
  then%0A%09%09
-then(
+Trace secondTrace = 
 traceKee
@@ -1575,17 +1575,99 @@
 race
-).isNull(
+;%0A%09%09then(secondTrace.getSpan().getTraceId()).isNotEqualTo(firstTrace.getSpan().getTraceId()
 );%0A%09
@@ -2121,16 +2121,21 @@
 ion %7B%0A%09%09
+this.
 executor
@@ -2160,16 +2160,21 @@
 unnable(
+this.
 traceMan
@@ -2314,24 +2314,29 @@
  run() %7B%0A%09%09%09
+this.
 trace = Trac
@@ -2498,74 +2498,13 @@
 %0A%09%09%09
-TraceContextHolder.setCurrentTrace(Mockito.mock(Trace.class));%0A%09%09%09
+this.
 trac

Fix @Ignored test