mockito_IAerospikeClient.scala

package com.travelaudience.data.job.tracking

import TravelAudience.TrackingProto.TAF.taf.TAFEvent
import com.aerospike.client.{IAerospikeClient, Key, Record}
import org.apache.beam.sdk.transforms.DoFn.ProcessElement
import org.apache.beam.sdk.transforms.{DoFn, DoFnTester}
import org.scalatest.{FlatSpec, Matchers}

import scala.collection.JavaConverters._
import scala.util.{Failure, Success, Try}

/* DoFn */

class MockitoDoFnThatInteractsWithAerospike(aerospikeClient: IAerospikeClient) extends DoFn[String, Int] {

  @ProcessElement
  def processElement(c: ProcessContext): Unit = {
    val uuid = c.element()

    // some interaction with aerospike...
    val triedRecord = Option(aerospikeClient.get(null, new Key("namespace", "set", uuid)))

    // do your things here that have to be done...
    triedRecord match {
      case Some(record) =>
        Try(TAFEvent.parseFrom(record.getValue("bin").asInstanceOf[Array[Byte]])) match {
          case Success(taf) =>
            // output result
            c.output(taf.timestamp.get)
          case Failure(ex) =>
        }
      case _ =>
    }
  }
}

/* Testing DoFn */

class MockitoDoFnThatInteractsWithAerospikeSpec extends FlatSpec with Matchers {

  import org.mockito.Mockito._

  val aerospikeClient: IAerospikeClient = mock(classOf[IAerospikeClient], withSettings().serializable())

  private val doFnTester = DoFnTester.of(new MockitoDoFnThatInteractsWithAerospike(aerospikeClient))

  "MockitoDoFnThatInteractsWithAerospike" should " work." in {

    //1. Set expectations about the normal Aerospike behaviour
    when(aerospikeClient.get(null, new Key("namespace", "set", "uuid1"))) thenReturn {
      val bin = Map("bin" -> TAFEvent(Some(123)).toByteArray.asInstanceOf[Object]).asJava
      val generation = 1
      val expiration = 100
      new Record(bin, generation, expiration)
    }

    when(aerospikeClient.get(null, new Key("namespace", "set", "uuid2"))) thenReturn {
      val bin = Map("bin" -> TAFEvent(Some(456)).toByteArray.asInstanceOf[Object]).asJava
      val generation = 1
      val expiration = 100
      new Record(bin, generation, expiration)
    }

    when(aerospikeClient.get(null, new Key("namespace", "set", "uuid3"))) thenReturn {
      val bin = Map("bin" -> TAFEvent(Some(789)).toByteArray.asInstanceOf[Object]).asJava
      val generation = 1
      val expiration = 100
      new Record(bin, generation, expiration)
    }

    // for any input
    // import org.mockito.ArgumentMatchers._
    // when(aerospikeClient.get(isNull, any(classOf[Key]))).thenReturn(record)

    // Test it
    val actualOutput = doFnTester.processBundle(Iterable("uuid1", "uuid3").asJava)
    val expectedOutput = List(123, 789).asJava

    // Check outputs
    actualOutput shouldEqual expectedOutput
  }

}

"org.mockito" % "mockito-core" % "2.13.0" % Test