/*
 * -\-\-
 * Spotify Styx Scheduler Service
 * --
 * Copyright (C) 2018 Spotify AB
 * --
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * -/-/-
 */

package com.spotify.styx.util;

import static com.spotify.styx.util.MDCUtil.withMDC;
import static java.util.Collections.emptyMap;
import static org.hamcrest.Matchers.anyOf;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ForkJoinPool;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.MDC;

public class MDCUtilTest {

  @Before
  public void setUp() throws Exception {
    MDC.clear();
  }

  @After
  public void tearDown() throws Exception {
    MDC.clear();
  }

  @Test
  public void withMDCRunnable() throws ExecutionException, InterruptedException {
    MDC.put("foo", "bar");
    Executors.newSingleThreadExecutor()
        .submit(withMDC(() -> assertThat(MDC.get("foo"), is("bar"))))
        .get();
  }

  @Test
  public void withMDCCallable() throws ExecutionException, InterruptedException {
    MDC.put("foo", "bar");
    Executors.newSingleThreadExecutor()
        .submit(withMDC(() -> {
          assertThat(MDC.get("foo"), is("bar"));
          return null;
        })).get();
  }

  @Test
  public void withMDCCommonPool() throws ExecutionException, InterruptedException {
    MDC.put("foo", "bar");
    CompletableFuture.runAsync(
        () -> assertThat(MDC.get("foo"), is("bar")),
        withMDC())
        // Later stages should also have the MDC applied
        .thenRun(() -> assertThat(MDC.get("foo"), is("bar")))
        .get();

    // MDC should not leak
    ForkJoinPool.commonPool()
        .submit(() -> assertThat(MDC.getCopyOfContextMap(), is(anyOf(nullValue(), is(emptyMap())))))
        .get();
  }

  @Test
  public void withMDCExecutor() throws ExecutionException, InterruptedException {
    MDC.put("foo", "bar");
    final ExecutorService executor = Executors.newSingleThreadExecutor();
    CompletableFuture.runAsync(
        () -> assertThat(MDC.get("foo"), is("bar")),
        withMDC(executor))
        // Later stages should also have the MDC applied
        .thenRun(() -> assertThat(MDC.get("foo"), is("bar")))
        .get();

    // MDC should not leak
    executor
        .submit(() -> assertThat(MDC.getCopyOfContextMap(), is(anyOf(nullValue(), is(emptyMap())))))
        .get();
  }
}

@@ -1305,16 +1305,104 @@
 Test %7B%0A%0A
+  private static final ExecutorService EXECUTOR = Executors.newSingleThreadExecutor();%0A%0A
   @Befor
@@ -1666,52 +1666,87 @@
     
-Executors.newSingleThreadExecutor
+final CompletableFuture%3CString%3E value = new CompletableFuture%3C%3E
 ()
+;
 %0A    
-    
+EXECUTOR
 .sub
@@ -1759,34 +1759,38 @@
 thMDC(() -%3E 
-assertThat
+value.complete
 (MDC.get(%22fo
@@ -1796,36 +1796,50 @@
 oo%22)
-, is(%22bar%22)
 )))
+;
 %0A    
-    .get(
+assertThat(value.get(), is(%22bar%22)
 );%0A 
@@ -1967,164 +1967,114 @@
     
-Executors.newSingleThreadExecutor()%0A        .submit(withMDC(() -%3E %7B%0A          assertThat(MDC.get(%22foo%22), is(%22bar%22));%0A          return null;%0A        %7D)).get(
+final String value = EXECUTOR.submit(withMDC(() -%3E MDC.get(%22foo%22))).get();%0A    assertThat(value, is(%22bar%22)
 );%0A 
@@ -2736,82 +2736,8 @@
 %22);%0A
-    final ExecutorService executor = Executors.newSingleThreadExecutor();%0A
     
@@ -2833,24 +2833,24 @@
 withMDC(
-executor
+EXECUTOR
 ))%0A     
@@ -3015,16 +3015,16 @@
     
-executor
+EXECUTOR
 %0A   

make codacy happy