org.mockito.stubbing.Answer Scala Examples

package com.basho.riak.spark.rdd.timeseries

import com.basho.riak.client.core.query.timeseries.ColumnDescription.ColumnType
import com.basho.riak.client.core.query.timeseries.{Cell, ColumnDescription}
import com.basho.riak.spark.query.{QueryTS, TSQueryData}
import com.basho.riak.spark.rdd.{RegressionTests, RiakTSRDD}
import com.basho.riak.spark.rdd.connector.RiakConnector
import com.basho.riak.spark.rdd.partitioner.RiakTSPartition
import org.apache.spark.{SparkContext, TaskContext}
import org.junit.Assert._
import org.junit.Test
import org.junit.runner.RunWith
import org.mockito.Matchers._
import org.mockito.Mockito._
import org.powermock.modules.junit4.PowerMockRunner
import org.powermock.api.mockito.PowerMockito
import com.basho.riak.client.core.query.timeseries.{Row => RiakRow}
import org.apache.spark.sql.{Row => SparkRow}
import org.junit.experimental.categories.Category
import org.mockito.Mock
import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.powermock.core.classloader.annotations.PrepareForTest

import scala.collection.JavaConversions._

@RunWith(classOf[PowerMockRunner])
@PrepareForTest(Array(classOf[QueryTS],classOf[RiakTSRDD[_]]))
class RiakTSRDDTest {

  @Mock
  protected val rc: RiakConnector = null

  @Mock
  protected val sc: SparkContext = null

  @Mock
  protected val tc: TaskContext = null

  @Test
  @Category(Array(classOf[RegressionTests]))
  def readAll(): Unit = {
    val rdd = new RiakTSRDD[SparkRow](sc, rc, "test")

    val neTsQD = TSQueryData("non-empty", None)
    val nonEmptyResponse = (Seq(new ColumnDescription("Col1", ColumnType.VARCHAR), new ColumnDescription("Col2", ColumnType.SINT64)),
      Seq(new RiakRow(List(new Cell("string-value"), new Cell(42)))))

    PowerMockito.whenNew(classOf[QueryTS]).withAnyArguments().thenAnswer( new Answer[QueryTS]{
      override def answer(invocation: InvocationOnMock): QueryTS = {
        val args = invocation.getArguments
        val rc: RiakConnector = args(0).asInstanceOf[RiakConnector]
        val qd: Seq[TSQueryData] = args(1).asInstanceOf[Seq[TSQueryData]]

        val q = spy(QueryTS(rc, qd))

        // By default returns an empty result
        doReturn(Seq() -> Seq()).when(q).nextChunk(any[TSQueryData])

        // Return 1 row for non-empty result
        doReturn(nonEmptyResponse).when(q).nextChunk(org.mockito.Matchers.eq(neTsQD))

        q
      }
    })

    // ----------  Perform test
    val iterator = rdd.compute( RiakTSPartition(0, Nil, List( null, null, null, neTsQD, neTsQD, null, neTsQD, null)), tc)

    // ----------  verify results
    val seq: Seq[SparkRow] = iterator.toIndexedSeq

    
    assertEquals(3, seq.size)

    seq.foreach(r => {
      assertEquals(2, r.size)
      assertEquals("string-value", r.get(0))
      assertEquals(42l, r.get(1))
    })
  }
} 

package me.snov.akka.sqs.shape

import akka.Done
import akka.stream.scaladsl.{Keep, Sink}
import akka.stream.testkit.scaladsl.TestSource
import com.amazonaws.handlers.AsyncHandler
import com.amazonaws.services.sqs.model._
import me.snov.akka.sqs._
import me.snov.akka.sqs.client.SqsClient
import org.mockito.Mockito._
import org.mockito.ArgumentMatchers._
import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.scalatest.mockito.MockitoSugar.mock
import org.scalatest.{FlatSpec, Matchers}

import scala.concurrent.Await
import scala.concurrent.duration._

class SqsAckSinkShapeSpec extends FlatSpec with Matchers with DefaultTestContext {
  it should "delete messages on Ack" in {

    val sqsClient = mock[SqsClient]
    when(sqsClient.deleteAsync(any(), any())).thenAnswer(
      new Answer[Object] {
        override def answer(invocation: InvocationOnMock): Object = {
          val receiptHandle = invocation.getArgument[String](0)
          val callback = invocation.getArgument[AsyncHandler[DeleteMessageRequest, DeleteMessageResult]](1)
          callback.onSuccess(
            new DeleteMessageRequest().withReceiptHandle(receiptHandle),
            new DeleteMessageResult
          )
          None
        }
      }
    )

    val (probe, future) = TestSource.probe[MessageActionPair]
      .toMat(Sink.fromGraph(SqsAckSinkShape(sqsClient)))(Keep.both)
      .run()

    probe
      .sendNext((new Message().withReceiptHandle("123"), Ack()))
      .sendComplete()

    Await.result(future, 1.second) shouldBe Done
    verify(sqsClient, times(1)).deleteAsync(any(), any())
  }

  it should "requeue messages on RequeueWithDelay" in {

    val sqsClient = mock[SqsClient]
    when(sqsClient.sendWithDelayAsync(any[String], any[Int], any())).thenAnswer(
      new Answer[Object] {
        override def answer(invocation: InvocationOnMock): Object = {
          val body = invocation.getArgument[String](0)
          val delay = invocation.getArgument[Int](1)
          val callback = invocation.getArgument[AsyncHandler[SendMessageRequest, SendMessageResult]](2)
          callback.onSuccess(
            new SendMessageRequest().withMessageBody(body).withDelaySeconds(delay),
            new SendMessageResult().withMessageId("12345")
          )
          None
        }
      }
    )

    val (probe, future) = TestSource.probe[MessageActionPair]
      .toMat(Sink.fromGraph(SqsAckSinkShape(sqsClient)))(Keep.both)
      .run()

    probe
      .sendNext((new Message().withBody("foo"), RequeueWithDelay(9)))
      .sendComplete()

    Await.result(future, 100.second) shouldBe Done
    verify(sqsClient, times(1)).sendWithDelayAsync(any(), any(), any())
  }
} 

package org.apache.toree.communication.socket

import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.scalatest.{Matchers, BeforeAndAfter, OneInstancePerTest, FunSpec}
import org.scalatest.mock.MockitoSugar
import org.mockito.Mockito._
import org.zeromq.ZMsg

class JeroMQSocketSpec extends FunSpec with MockitoSugar
  with OneInstancePerTest with BeforeAndAfter with Matchers
{
  private val runnable = mock[ZeroMQSocketRunnable]
  @volatile private var running = true
  //  Mock the running of the runnable for the tests
  doAnswer(new Answer[Unit] {
    override def answer(invocation: InvocationOnMock): Unit = while (running) {
      Thread.sleep(1)
    }
  }).when(runnable).run()


  //  Mock the close of the runnable to shutdown
  doAnswer(new Answer[Unit] {
    override def answer(invocation: InvocationOnMock): Unit = running = false
  }).when(runnable).close()

  private val socket: JeroMQSocket = new JeroMQSocket(runnable)

  after {
    running = false
  }

  describe("JeroMQSocket") {
    describe("#send") {
      it("should offer a message to the runnable") {
        val message: String = "Some Message"
        val expected = ZMsg.newStringMsg(message)

        socket.send(message.getBytes)
        verify(runnable).offer(expected)
      }

      it("should thrown and AssertionError when socket is no longer alive") {
        socket.close()

        intercept[AssertionError] {
          socket.send("".getBytes)
        }
      }
    }

    describe("#close") {
      it("should close the runnable") {
        socket.close()

        verify(runnable).close()
      }

      it("should close the socket thread") {
        socket.close()

        socket.isAlive should be (false)
      }
    }

    describe("#isAlive") {
      it("should evaluate to true when the socket thread is alive") {
        socket.isAlive should be (true)
      }

      it("should evaluate to false when the socket thread is dead") {
        socket.close()

        socket.isAlive should be (false)
      }
    }
  }
} 

package com.programmaticallyspeaking.ncd.transpile

import com.programmaticallyspeaking.ncd.testing.UnitTest
import org.scalatest.mockito.MockitoSugar
import org.mockito.Mockito._
import org.mockito.ArgumentMatchers._
import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer

import scala.collection.mutable.ListBuffer

class CachingES5TranspilerTest extends UnitTest with MockitoSugar {

  val transpilerCalledWith = ListBuffer[String]()

  def fakeTranspiler: ES5Transpiler = {
    transpilerCalledWith.clear()
    val t = mock[ES5Transpiler]
    when(t.transpile(any[String])).thenAnswer(new Answer[String] {
      override def answer(invocation: InvocationOnMock): String = {
        val input = invocation.getArgument[String](0)
        transpilerCalledWith += input
        "transpiled:" + input
      }
    })
    t
  }

  def createSut = new CachingES5Transpiler(fakeTranspiler)

  "A caching ES5 transpiler" - {
    "delegates to the real one" in {
      val caching = createSut
      caching.transpile("input") should be ("transpiled:input")
    }

    "delegates for each unknown input" in {
      val caching = createSut
      caching.transpile("input1")
      caching.transpile("input2")
      transpilerCalledWith should be (Seq("input1", "input2"))
    }

    "caches and reuses ouput" in {
      val caching = createSut
      caching.transpile("input1")
      caching.transpile("input1")
      transpilerCalledWith should be (Seq("input1"))
    }
  }

} 

package ai.deepsense.workflowmanager.rest

import scala.concurrent.Future

import org.mockito.Matchers._
import org.mockito.Mockito._
import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.scalatest.mockito.MockitoSugar
import spray.routing.Route

import ai.deepsense.commons.auth.usercontext.{CannotGetUserException, Role, TokenTranslator, UserContext}

trait ApiSpecSupport extends MockitoSugar {

  val authTokens: Map[String, Set[String]]

  val apiPrefix: String

  def createRestComponent(tokenTranslator: TokenTranslator): Route

  protected def testRoute: Route = {
    val tokenTranslator = mock[TokenTranslator]
    when(tokenTranslator.translate(any(classOf[String])))
      .thenAnswer(new Answer[Future[UserContext]] {
      override def answer(invocation: InvocationOnMock): Future[UserContext] = {
        val tokenFromRequest = invocation.getArgumentAt(0, classOf[String])
        if (authTokens.keySet.contains(tokenFromRequest)) {
          val uc = mockUserContext(tokenFromRequest)
          Future.successful(uc)
        } else {
          Future.failed(new CannotGetUserException(tokenFromRequest))
        }
      }
    })
    createRestComponent(tokenTranslator)
  }

  private def mockUserContext(tenantId: String): UserContext = {
    val userContext = mock[UserContext]
    when(userContext.tenantId).thenReturn(tenantId)
    when(userContext.roles).thenReturn(authTokens(tenantId).map(Role(_)))
    userContext
  }
} 

package io.hydrosphere.mist.core

import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.{Answer, OngoingStubbing}

import scala.concurrent.Future
import scala.reflect.ClassTag

trait MockitoSugar extends org.scalatest.mockito.MockitoSugar {

  def any[T <: AnyRef](implicit classTag: ClassTag[T]): T = {
    org.mockito.Matchers.any(classTag.runtimeClass.asInstanceOf[Class[T]])
  }

  def when[T](call: T): org.mockito.stubbing.OngoingStubbing[T] = {
    org.mockito.Mockito.when(call)
  }

  implicit class FutureStubbing[A](stubbing: OngoingStubbing[Future[A]]) {

    def thenSuccess(a: A): OngoingStubbing[Future[A]] =
      stubbing.thenReturn(Future.successful(a))

    def thenFailure(e: Throwable): OngoingStubbing[Future[A]] =
      stubbing.thenReturn(Future.failed(e))
  }

  implicit class RespondFuncSyntax[A](stubbing: OngoingStubbing[A]) {

    def thenRespond[A1](f: A1 => A): OngoingStubbing[A] = {
      stubbing.thenAnswer(new Answer[A] {
        def answer(invocation: InvocationOnMock): A = {
          val args = invocation.getArguments
          val a1 = args(0).asInstanceOf[A1]
          f(a1)
        }
      })
    }

    def thenRespond[A1, A2](f: (A1, A2) => A): OngoingStubbing[A] = {
      stubbing.thenAnswer(new Answer[A] {
        def answer(invocation: InvocationOnMock): A = {
          val args = invocation.getArguments
          val a1 = args(0).asInstanceOf[A1]
          val a2 = args(1).asInstanceOf[A2]
          f(a1, a2)
        }
      })
    }

    def thenRespond[A1, A2, A3](f: (A1, A2, A3) => A): OngoingStubbing[A] = {
      stubbing.thenAnswer(new Answer[A] {
        def answer(invocation: InvocationOnMock): A = {
          val args = invocation.getArguments
          val a1 = args(0).asInstanceOf[A1]
          val a2 = args(1).asInstanceOf[A2]
          val a3 = args(1).asInstanceOf[A3]
          f(a1, a2, a3)
        }
      })
    }
  }
} 

package org.apache.spark.deploy.k8s.submit.steps.initcontainer

import scala.collection.JavaConverters._

import io.fabric8.kubernetes.api.model._
import org.mockito.{Mock, MockitoAnnotations}
import org.mockito.Matchers.any
import org.mockito.Mockito.when
import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.scalatest.BeforeAndAfter

import org.apache.spark.SparkFunSuite
import org.apache.spark.deploy.k8s.{InitContainerBootstrap, PodWithDetachedInitContainer}
import org.apache.spark.deploy.k8s.Config._

class BasicInitContainerConfigurationStepSuite extends SparkFunSuite with BeforeAndAfter {

  private val SPARK_JARS = Seq(
    "hdfs://localhost:9000/app/jars/jar1.jar", "file:///app/jars/jar2.jar")
  private val SPARK_FILES = Seq(
    "hdfs://localhost:9000/app/files/file1.txt", "file:///app/files/file2.txt")
  private val JARS_DOWNLOAD_PATH = "/var/data/jars"
  private val FILES_DOWNLOAD_PATH = "/var/data/files"
  private val POD_LABEL = Map("bootstrap" -> "true")
  private val INIT_CONTAINER_NAME = "init-container"
  private val DRIVER_CONTAINER_NAME = "driver-container"

  @Mock
  private var podAndInitContainerBootstrap : InitContainerBootstrap = _

  before {
    MockitoAnnotations.initMocks(this)
    when(podAndInitContainerBootstrap.bootstrapInitContainer(
      any[PodWithDetachedInitContainer])).thenAnswer(new Answer[PodWithDetachedInitContainer] {
      override def answer(invocation: InvocationOnMock) : PodWithDetachedInitContainer = {
        val pod = invocation.getArgumentAt(0, classOf[PodWithDetachedInitContainer])
        pod.copy(
          pod = new PodBuilder(pod.pod)
            .withNewMetadata()
            .addToLabels("bootstrap", "true")
            .endMetadata()
            .withNewSpec().endSpec()
            .build(),
          initContainer = new ContainerBuilder()
            .withName(INIT_CONTAINER_NAME)
            .build(),
          mainContainer = new ContainerBuilder()
            .withName(DRIVER_CONTAINER_NAME)
            .build()
        )}})
  }

  test("additionalDriverSparkConf with mix of remote files and jars") {
    val baseInitStep = new BasicInitContainerConfigurationStep(
      SPARK_JARS,
      SPARK_FILES,
      JARS_DOWNLOAD_PATH,
      FILES_DOWNLOAD_PATH,
      podAndInitContainerBootstrap)
    val expectedDriverSparkConf = Map(
      JARS_DOWNLOAD_LOCATION.key -> JARS_DOWNLOAD_PATH,
      FILES_DOWNLOAD_LOCATION.key -> FILES_DOWNLOAD_PATH,
      INIT_CONTAINER_REMOTE_JARS.key -> "hdfs://localhost:9000/app/jars/jar1.jar",
      INIT_CONTAINER_REMOTE_FILES.key -> "hdfs://localhost:9000/app/files/file1.txt")
    val initContainerSpec = InitContainerSpec(
      Map.empty[String, String],
      Map.empty[String, String],
      new Container(),
      new Container(),
      new Pod,
      Seq.empty[HasMetadata])
    val returnContainerSpec = baseInitStep.configureInitContainer(initContainerSpec)
    assert(expectedDriverSparkConf === returnContainerSpec.properties)
    assert(returnContainerSpec.initContainer.getName === INIT_CONTAINER_NAME)
    assert(returnContainerSpec.driverContainer.getName === DRIVER_CONTAINER_NAME)
    assert(returnContainerSpec.driverPod.getMetadata.getLabels.asScala === POD_LABEL)
  }
} 

package dcos.metronome
package utils.test

import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.{Answer, OngoingStubbing}
import org.mockito.verification.VerificationMode
import org.mockito.{ArgumentMatchers, Mockito => M}
import org.scalatestplus.mockito.MockitoSugar


trait Mockito extends MockitoSugar {

  def eq[T](t: T) = ArgumentMatchers.eq(t)
  def any[T] = ArgumentMatchers.any[T]
  def anyBoolean = ArgumentMatchers.anyBoolean
  def anyString = ArgumentMatchers.anyString
  def same[T](value: T) = ArgumentMatchers.same(value)
  def verify[T](t: T, mode: VerificationMode = times(1)) = M.verify(t, mode)
  def times(num: Int) = M.times(num)
  def timeout(millis: Int) = M.timeout(millis.toLong)
  def atLeastOnce = M.atLeastOnce()
  def atLeast(num: Int) = M.atLeast(num)
  def atMost(num: Int) = M.atMost(num)
  def never = M.never()

  def inOrder(mocks: AnyRef*) = M.inOrder(mocks: _*)

  def noMoreInteractions(mocks: AnyRef*): Unit = {
    M.verifyNoMoreInteractions(mocks: _*)
  }

  def reset(mocks: AnyRef*): Unit = {
    M.reset(mocks: _*)
  }

  class MockAnswer[T](function: Array[AnyRef] => T) extends Answer[T] {
    def answer(invocation: InvocationOnMock): T = {
      function(invocation.getArguments)
    }
  }

  implicit class Stubbed[T](c: => T) {
    def returns(t: T, t2: T*): OngoingStubbing[T] = {
      if (t2.isEmpty) M.when(c).thenReturn(t)
      else t2.foldLeft(M.when(c).thenReturn(t)) { (res, cur) => res.thenReturn(cur) }
    }
    def answers(function: Array[AnyRef] => T) = M.when(c).thenAnswer(new MockAnswer(function))
    def throws[E <: Throwable](e: E*): OngoingStubbing[T] = {
      if (e.isEmpty) throw new java.lang.IllegalArgumentException("The parameter passed to throws must not be empty")
      e.drop(1).foldLeft(M.when(c).thenThrow(e.head)) { (res, cur) => res.thenThrow(cur) }
    }
  }
}

object Mockito extends Mockito 

package org.apache.toree.communication.socket

import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.scalatest.{Matchers, BeforeAndAfter, OneInstancePerTest, FunSpec}
import org.scalatest.mock.MockitoSugar
import org.mockito.Mockito._
import org.zeromq.ZMsg

class JeroMQSocketSpec extends FunSpec with MockitoSugar
  with OneInstancePerTest with BeforeAndAfter with Matchers
{
  private val runnable = mock[ZeroMQSocketRunnable]
  @volatile private var running = true
  //  Mock the running of the runnable for the tests
  doAnswer(new Answer[Unit] {
    override def answer(invocation: InvocationOnMock): Unit = while (running) {
      Thread.sleep(1)
    }
  }).when(runnable).run()


  //  Mock the close of the runnable to shutdown
  doAnswer(new Answer[Unit] {
    override def answer(invocation: InvocationOnMock): Unit = running = false
  }).when(runnable).close()

  private val socket: JeroMQSocket = new JeroMQSocket(runnable)

  after {
    running = false
  }

  describe("JeroMQSocket") {
    describe("#send") {
      it("should offer a message to the runnable") {
        val message: String = "Some Message"
        val expected = ZMsg.newStringMsg(message)

        socket.send(message.getBytes)
        verify(runnable).offer(expected)
      }

      it("should thrown and AssertionError when socket is no longer alive") {
        socket.close()

        intercept[AssertionError] {
          socket.send("".getBytes)
        }
      }
    }

    describe("#close") {
      it("should close the runnable") {
        socket.close()

        verify(runnable).close()
      }

      it("should close the socket thread") {
        socket.close()

        socket.isAlive should be (false)
      }
    }

    describe("#isAlive") {
      it("should evaluate to true when the socket thread is alive") {
        socket.isAlive should be (true)
      }

      it("should evaluate to false when the socket thread is dead") {
        socket.close()

        socket.isAlive should be (false)
      }
    }
  }
}