Fix MdcLoggingExecutionContext and tests

http-verbs-common/src/main/scala/uk/gov/hmrc/play/http/logging/MdcLoggingExecutionContext.scala

@@ -29,13 +29,18 @@ class MdcLoggingExecutionContext(wrapped: ExecutionContext, mdcData: Map[String,
 
   private class RunWithMDC(runnable: Runnable, mdcData: Map[String, String]) extends Runnable {
     def run(): Unit = {
+      val oldMdcData = MDC.getCopyOfContextMap
+      MDC.clear()
       mdcData.foreach {
         case (k, v) => MDC.put(k, v)
       }
       try {
         runnable.run()
       } finally {
-        MDC.clear()
+        if (oldMdcData == null)
+          MDC.clear()
+        else
+          MDC.setContextMap(oldMdcData)
       }
     }
   }

http-verbs-common/src/test/scala/uk/gov/hmrc/http/RetriesSpec.scala

@@ -193,6 +193,8 @@ class RetriesSpec
 
       val expectedResponse = HttpResponse(404, "")
 
+      org.slf4j.MDC.clear()
+
       val resultF =
         for {
           _   <- Future.successful(Mdc.putMdc(mdcData))

http-verbs-common/src/test/scala/uk/gov/hmrc/http/client/HttpClientV2Spec.scala

@@ -603,6 +603,8 @@ class HttpClientV2Spec
 
       implicit val mdcEc = ExecutionContext.fromExecutor(new uk.gov.hmrc.play.http.logging.MDCPropagatingExecutorService(Executors.newFixedThreadPool(2)))
 
+      org.slf4j.MDC.clear()
+
       val res: Future[Map[String, String]] =
         for {
           _ <- Future.successful(Mdc.putMdc(mdcData))

http-verbs-common/src/test/scala/uk/gov/hmrc/play/http/logging/MdcLoggingExecutionContextSpec.scala

@@ -27,28 +27,25 @@ import org.scalatest.matchers.should.Matchers
 import org.slf4j.{LoggerFactory, MDC}
 import play.core.NamedThreadFactory
 
-import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.concurrent.duration._
 import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.reflect._
 
 @annotation.nowarn("msg=deprecated")
 class MdcLoggingExecutionContextSpec
-    extends AnyWordSpecLike
-    with Matchers
-    with LoneElement
-    with Inspectors
-    with BeforeAndAfter {
+  extends AnyWordSpecLike
+     with Matchers
+     with LoneElement
+     with Inspectors
+     with BeforeAndAfter {
 
   before {
     MDC.clear()
   }
 
   "The MDC Transporting Execution Context" should {
-
-    "capture the an MDC map with values in it and put it in place when a task is run" in withCaptureOfLoggingFrom[
-      MdcLoggingExecutionContextSpec] { logList =>
+    "capture the MDC map with values in it and put it in place when a task is run" in withCaptureOfLoggingFrom[MdcLoggingExecutionContextSpec] { logList =>
       implicit val ec = createAndInitialiseMdcTransportingExecutionContext(Map("someKey" -> "something"))
 
       logEventInsideAFutureUsing(ec)
@@ -75,20 +72,18 @@ class MdcLoggingExecutionContextSpec
       logList.loneElement._2 should be(Map("someKey" -> "something"))
     }
 
-    "clear the MDC map after a task throws an exception" in withCaptureOfLoggingFrom[MdcLoggingExecutionContextSpec] {
-      logList =>
-        implicit val ec = createAndInitialiseMdcTransportingExecutionContext(Map("someKey" -> "something"))
+    "clear the MDC map after a task throws an exception" in withCaptureOfLoggingFrom[MdcLoggingExecutionContextSpec] { logList =>
+      implicit val ec = createAndInitialiseMdcTransportingExecutionContext(Map("someKey" -> "something"))
 
-        throwAnExceptionInATaskOn(ec)
+      throwAnExceptionInATaskOn(ec)
 
-        MDC.clear()
-        logEventInsideAFutureUsing(ec)
+      MDC.clear()
+      logEventInsideAFutureUsing(ec)
 
-        logList.loneElement._2 should be(Map("someKey" -> "something"))
+      logList.loneElement._2 should be(Map("someKey" -> "something"))
     }
 
-    "log values from given MDC map when multiple threads are using it concurrently by ensuring each log from each thread has been logged via MDC" in withCaptureOfLoggingFrom[
-      MdcLoggingExecutionContextSpec] { logList =>
+    "log values from given MDC map when multiple threads are using it concurrently by ensuring each log from each thread has been logged via MDC" in withCaptureOfLoggingFrom[MdcLoggingExecutionContextSpec] { logList =>
       val threadCount = 10
       val logCount    = 10
 
@@ -131,36 +126,33 @@ class MdcLoggingExecutionContextSpec
     ec
   }
 
-  def logEventInsideAFutureUsingImplicitEc(implicit ec: ExecutionContext) {
+  def logEventInsideAFutureUsingImplicitEc(implicit ec: ExecutionContext): Unit =
     logEventInsideAFutureUsing(ec)
-  }
 
-  def logEventInsideAFutureUsing(ec: ExecutionContext) {
-    Await.ready(Future {
-      LoggerFactory.getLogger(classOf[MdcLoggingExecutionContextSpec]).info("")
-    }(ec), 2.second)
-  }
+  def logEventInsideAFutureUsing(ec: ExecutionContext): Unit =
+    Await.ready(
+      Future.apply(
+        LoggerFactory.getLogger(classOf[MdcLoggingExecutionContextSpec]).info("")
+      )(ec),
+      2.second
+    )
 
-  def doSomethingInsideAFutureButDontLog(ec: ExecutionContext) {
-    Await.ready(Future {}(ec), 2.second)
-  }
+  def doSomethingInsideAFutureButDontLog(ec: ExecutionContext): Unit =
+    Await.ready(Future.apply()(ec), 2.second)
 
-  def throwAnExceptionInATaskOn(ec: ExecutionContext) {
-    ec.execute(new Runnable() {
-      def run(): Unit =
-        throw new RuntimeException("Test what happens when a task running on this EC throws an exception")
-    })
-  }
+  def throwAnExceptionInATaskOn(ec: ExecutionContext): Unit =
+    ec.execute(() => throw new RuntimeException("Test what happens when a task running on this EC throws an exception"))
 
   /** Ensures that a thread is already created in the execution context by running an empty future.
     * Required as otherwise the MDC is transferred to the new thread as it is stored in an inheritable
     * ThreadLocal.
     */
-  def initialise(ec: ExecutionContext) {
-    Await.ready(Future {}(ec), 2.second)
-  }
+  def initialise(ec: ExecutionContext): Unit =
+    Await.ready(Future.apply()(ec), 2.second)
+
+  def withCaptureOfLoggingFrom[T: ClassTag](body: (=> List[(ILoggingEvent, Map[String, String])]) => Unit): Unit = {
+    import scala.collection.JavaConverters._
 
-  def withCaptureOfLoggingFrom[T: ClassTag](body: (=> List[(ILoggingEvent, Map[String, String])]) => Any) = {
     val logger = LoggerFactory.getLogger(classTag[T].runtimeClass).asInstanceOf[LogbackLogger]
     val appender = new AppenderBase[ILoggingEvent]() {
 

http-verbs-common/src/test/scala/uk/gov/hmrc/play/http/logging/MdcSpec.scala

@@ -18,7 +18,7 @@ package uk.gov.hmrc.play.http.logging
 
 import akka.dispatch.ExecutorServiceDelegate
 import org.scalatest.BeforeAndAfter
-import org.scalatest.concurrent.ScalaFutures
+import org.scalatest.concurrent.{IntegrationPatience, ScalaFutures}
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.wordspec.AnyWordSpecLike
 import org.slf4j.MDC
@@ -31,6 +31,7 @@ class MdcSpec
   extends AnyWordSpecLike
      with Matchers
      with ScalaFutures
+     with IntegrationPatience
      with BeforeAndAfter {
 
   before {
@@ -96,19 +97,16 @@ class MdcSpec
 class MDCPropagatingExecutorService(val executor: ExecutorService) extends ExecutorServiceDelegate {
 
   override def execute(command: Runnable): Unit = {
-
     val mdcData = MDC.getCopyOfContextMap
 
-    executor.execute(new Runnable {
-      override def run(): Unit = {
-        val oldMdcData = MDC.getCopyOfContextMap
-        setMDC(mdcData)
-        try {
-          command.run()
-        } finally {
-          // this means any Mdc updates on the ec will not be propagated once it steps out
-          setMDC(oldMdcData)
-        }
+    executor.execute(() => {
+      val oldMdcData = MDC.getCopyOfContextMap
+      setMDC(mdcData)
+      try {
+        command.run()
+      } finally {
+        // this means any Mdc updates on the ec will not be propagated once it steps out
+        setMDC(oldMdcData)
       }
     })
   }