Make object mocking thread safe (#312)

Protect against concurrent non-mocked access to the a mocked object (Fix #311)

Author: ultrasecreth

build.sbt

@@ -43,7 +43,15 @@ lazy val commonSettings =
           Nil
       }
     },
-    Test / scalacOptions += "-Ywarn-value-discard"
+    Test / scalacOptions += "-Ywarn-value-discard",
+    libraryDependencies ++= {
+      CrossVersion.partialVersion(scalaVersion.value) match {
+        case Some((2, major)) if major <= 12 =>
+          Seq()
+        case _ =>
+          Seq("org.scala-lang.modules" %% "scala-parallel-collections" % "1.0.0-RC1")
+      }
+    }
   )
 
 lazy val publishSettings = Seq(

common/src/main/scala/org/mockito/MockitoAPI.scala

@@ -16,13 +16,13 @@ import org.mockito.ReflectionUtils.InvocationOnMockOps
 import org.mockito.internal.configuration.plugins.Plugins.getMockMaker
 import org.mockito.internal.creation.MockSettingsImpl
 import org.mockito.internal.exceptions.Reporter.notAMockPassedToVerifyNoMoreInteractions
-import org.mockito.internal.handler.ScalaMockHandler
+import org.mockito.internal.handler.{ ScalaMockHandler, ThreadAwareMockHandler }
 import org.mockito.internal.progress.ThreadSafeMockingProgress.mockingProgress
 import org.mockito.internal.stubbing.answers.ScalaThrowsException
 import org.mockito.internal.util.MockUtil
 import org.mockito.internal.util.reflection.LenientCopyTool
 import org.mockito.internal.{ ValueClassExtractor, ValueClassWrapper }
-import org.mockito.invocation.InvocationOnMock
+import org.mockito.invocation.{ Invocation, InvocationContainer, InvocationOnMock, MockHandler }
 import org.mockito.mock.MockCreationSettings
 import org.mockito.stubbing._
 import org.mockito.verification.{ VerificationAfterDelay, VerificationMode, VerificationWithTimeout }
@@ -472,7 +472,8 @@ private[mockito] trait MockitoEnhancer extends MockCreator {
    * <code>verify(aMock).iHaveSomeDefaultArguments("I'm not gonna pass the second argument", "default value")</code>
    * as the value for the second parameter would have been null...
    */
-  override def mock[T <: AnyRef: ClassTag: WeakTypeTag](implicit defaultAnswer: DefaultAnswer, $pt: Prettifier): T = mock(withSettings)
+  override def mock[T <: AnyRef: ClassTag: WeakTypeTag](implicit defaultAnswer: DefaultAnswer, $pt: Prettifier): T =
+    createMock(withSettings)
 
   /**
    * Delegates to <code>Mockito.mock(type: Class[T], defaultAnswer: Answer[_])</code>
@@ -489,7 +490,7 @@ private[mockito] trait MockitoEnhancer extends MockCreator {
    * as the value for the second parameter would have been null...
    */
   override def mock[T <: AnyRef: ClassTag: WeakTypeTag](defaultAnswer: DefaultAnswer)(implicit $pt: Prettifier): T =
-    mock(withSettings(defaultAnswer))
+    createMock(withSettings(defaultAnswer))
 
   /**
    * Delegates to <code>Mockito.mock(type: Class[T], mockSettings: MockSettings)</code>
@@ -505,7 +506,13 @@ private[mockito] trait MockitoEnhancer extends MockCreator {
    * <code>verify(aMock).iHaveSomeDefaultArguments("I'm not gonna pass the second argument", "default value")</code>
    * as the value for the second parameter would have been null...
    */
-  override def mock[T <: AnyRef: ClassTag: WeakTypeTag](mockSettings: MockSettings)(implicit $pt: Prettifier): T = {
+  override def mock[T <: AnyRef: ClassTag: WeakTypeTag](mockSettings: MockSettings)(implicit $pt: Prettifier): T =
+    createMock(mockSettings)
+
+  private def createMock[T <: AnyRef: ClassTag: WeakTypeTag](
+      mockSettings: MockSettings,
+      mockHandler: (MockCreationSettings[T], Prettifier) => MockHandler[T] = (settings: MockCreationSettings[T], pt: Prettifier) => ScalaMockHandler(settings)(pt)
+  )(implicit $pt: Prettifier): T = {
     val interfaces = ReflectionUtils.extraInterfaces
 
     val realClass: Class[T] = mockSettings match {
@@ -520,7 +527,7 @@ private[mockito] trait MockitoEnhancer extends MockCreator {
       else mockSettings
 
     def createMock(settings: MockCreationSettings[T]): T = {
-      val mock          = getMockMaker.createMock(settings, ScalaMockHandler(settings))
+      val mock          = getMockMaker.createMock(settings, mockHandler(settings, $pt))
       val spiedInstance = settings.getSpiedInstance
       if (spiedInstance != null) new LenientCopyTool().copyToMock(spiedInstance, mock)
       mock
@@ -620,12 +627,21 @@ private[mockito] trait MockitoEnhancer extends MockCreator {
   /**
    * Mocks the specified object only for the context of the block
    */
-  def withObjectMocked[O <: AnyRef: ClassTag](block: => Any): Unit = {
-    val moduleField = clazz[O].getDeclaredField("MODULE$")
-    val realImpl    = moduleField.get(null)
-    ReflectionUtils.setFinalStatic(moduleField, mock[O])
-    try block
-    finally ReflectionUtils.setFinalStatic(moduleField, realImpl)
+  def withObjectMocked[O <: AnyRef: ClassTag](block: => Any)(implicit defaultAnswer: DefaultAnswer, $pt: Prettifier): Unit = {
+    val objectClass = clazz[O]
+    objectClass.synchronized {
+      val moduleField = objectClass.getDeclaredField("MODULE$")
+      val realImpl: O = moduleField.get(null).asInstanceOf[O]
+
+      val threadAwareMock = createMock(
+        withSettings(defaultAnswer),
+        (settings: MockCreationSettings[O], pt: Prettifier) => ThreadAwareMockHandler(settings, realImpl)(pt)
+      )
+
+      ReflectionUtils.setFinalStatic(moduleField, threadAwareMock)
+      try block
+      finally ReflectionUtils.setFinalStatic(moduleField, realImpl)
+    }
   }
 }
 

common/src/main/scala/org/mockito/ReflectionUtils.scala

@@ -123,7 +123,7 @@ object ReflectionUtils {
         .getOrElse(Seq.empty)
     }
 
-  def setFinalStatic(field: Field, newValue: Any) = {
+  def setFinalStatic(field: Field, newValue: Any): Unit = {
     field.setAccessible(true)
     val modifiersField = classOf[Field].getDeclaredField("modifiers")
     modifiersField.setAccessible(true)

common/src/main/scala/org/mockito/internal/handler/ThreadAwareMockHandler.scala

@@ -0,0 +1,25 @@
+package org.mockito.internal.handler
+
+import org.mockito.AdditionalAnswers
+import org.mockito.invocation.{ Invocation, InvocationContainer, MockHandler }
+import org.mockito.mock.MockCreationSettings
+import org.scalactic.Prettifier
+
+class ThreadAwareMockHandler[T](settings: MockCreationSettings[T], realImpl: T)(implicit $pt: Prettifier) extends MockHandler[T] {
+  private val currentThread    = Thread.currentThread()
+  private val mockDelegate     = ScalaMockHandler(settings)
+  private val realImplDelegate = AdditionalAnswers.delegatesTo(realImpl)
+
+  override def handle(invocation: Invocation): AnyRef =
+    if (Thread.currentThread() == currentThread) mockDelegate.handle(invocation)
+    else realImplDelegate.answer(invocation)
+
+  override def getMockSettings: MockCreationSettings[T] = mockDelegate.getMockSettings
+
+  override def getInvocationContainer: InvocationContainer = mockDelegate.getInvocationContainer
+}
+
+object ThreadAwareMockHandler {
+  def apply[T](settings: MockCreationSettings[T], realImpl: T)(implicit $pt: Prettifier): ThreadAwareMockHandler[T] =
+    new ThreadAwareMockHandler(settings, realImpl)($pt)
+}

scalatest/src/test/scala/user/org/mockito/IdiomaticStubbingTest.scala

@@ -1,15 +1,16 @@
 package user.org.mockito
 
-import java.lang.reflect.{ Field, Modifier }
 import java.util.concurrent.atomic.AtomicInteger
 
 import org.mockito.invocation.InvocationOnMock
-import org.mockito.{ clazz, ArgumentMatchersSugar, IdiomaticStubbing }
+import org.mockito.{ ArgumentMatchersSugar, IdiomaticStubbing }
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.wordspec.AnyWordSpec
 import user.org.mockito.matchers.{ ValueCaseClassInt, ValueCaseClassString, ValueClass }
 
-import scala.reflect.ClassTag
+import scala.collection.parallel.immutable
+import scala.concurrent.{ Await, Future }
+import scala.util.Random
 
 class IdiomaticStubbingTest extends AnyWordSpec with Matchers with ArgumentMatchersSugar with IdiomaticMockitoTestSetup with IdiomaticStubbing {
 
@@ -313,5 +314,26 @@ class IdiomaticStubbingTest extends AnyWordSpec with Matchers with ArgumentMatch
 
       FooObject.simpleMethod shouldBe "not mocked!"
     }
+
+    "object stubbing should be thread safe" in {
+      immutable.ParSeq.range(1, 100).foreach { i =>
+        withObjectMocked[FooObject.type] {
+          FooObject.simpleMethod returns s"mocked!-$i"
+          FooObject.simpleMethod shouldBe s"mocked!-$i"
+        }
+      }
+    }
+
+    "object stubbing should be thread safe 2" in {
+      val now = FooObject.stateDependantMethod
+      immutable.ParSeq.range(1, 100).foreach { i =>
+        if (i % 2 == 0)
+          withObjectMocked[FooObject.type] {
+            FooObject.stateDependantMethod returns i
+            FooObject.stateDependantMethod shouldBe i
+          }
+        else FooObject.stateDependantMethod shouldBe now
+      }
+    }
   }
 }

scalatest/src/test/scala/user/org/mockito/TestModel.scala

@@ -2,6 +2,7 @@ package user.org.mockito
 import user.org.mockito.matchers.{ ValueCaseClassInt, ValueCaseClassString, ValueClass }
 
 import scala.annotation.varargs
+import scala.util.Random
 
 trait FooTrait {
   def bar = "not mocked"
@@ -128,5 +129,9 @@ class TestController(org: Org) {
 }
 
 object FooObject {
+  val now: Long = Random.nextLong()
+
   def simpleMethod: String = "not mocked!"
+
+  def stateDependantMethod: Long = now
 }