tests/src/test/scala/system/utils/KafkaUtils.scala - openwhisk-package-kafka - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You 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 system.utils

 import java.util.HashMap
 import java.util.Properties
 import java.util.concurrent.{TimeUnit, TimeoutException}

 import io.restassured.RestAssured
 import io.restassured.config.{RestAssuredConfig, SSLConfig}
 import javax.security.auth.login.Configuration
 import javax.security.auth.login.AppConfigurationEntry
 import org.apache.kafka.clients.producer.KafkaProducer

 import scala.collection.mutable.ListBuffer
 import spray.json.DefaultJsonProtocol._
 import spray.json._
 import system.packages.ActionHelper._
 import org.apache.openwhisk.utils.JsHelpers

 import scala.concurrent.duration.DurationInt
 import common.TestHelpers
 import common.TestUtils
 import common.WskTestHelpers
 import common.ActivationResult
 import org.apache.openwhisk.utils.retry
 import org.apache.kafka.clients.producer.ProducerRecord

 trait KafkaUtils extends TestHelpers with WskTestHelpers {
     lazy val messageHubProps = KafkaUtils.initializeMessageHub()

     def createProducer() : KafkaProducer[String, String] = {
         // currently only supporting MH
         new KafkaProducer[String, String](KafkaUtils.asKafkaProducerProps(this.messageHubProps))
     }

     def apply(key : String) = {
         this.messageHubProps.getOrElse(key, "")
     }

     def getAsJson(key : String) = {
         key match {
             case key if key == "brokers" => this(key).asInstanceOf[List[String]].toJson
             case key => this(key).asInstanceOf[String].toJson
         }
     }

     val sslconfig = {
         val inner = new SSLConfig().allowAllHostnames()
         val config = inner.relaxedHTTPSValidation()
         new RestAssuredConfig().sslConfig(config)
     }

     def createTrigger(assetHelper: AssetCleaner, name: String, parameters: Map[String, spray.json.JsValue]): String = {
         println(s"Creating trigger $name")

         val feedCreationResult = assetHelper.withCleaner(wsk.trigger, name) {
             (trigger, _) =>
                 trigger.create(name, feed = Some(s"/whisk.system/messaging/messageHubFeed"), parameters = parameters)
         }

         val activation = wsk.parseJsonString(feedCreationResult.stdout.substring(0, feedCreationResult.stdout.indexOf("ok: created trigger"))).convertTo[ActivationResult]

         // should be successful
         activation.response.success shouldBe true

         // It takes a moment for the consumer to fully initialize.
         println("Giving the consumer a moment to get ready")
         Thread.sleep(KafkaUtils.consumerInitTime)

         val uuid = activation.response.result.get.fields.get("uuid").get.toString().replaceAll("\"", "")
         consumerExists(uuid)

         uuid
     }


     def consumerExists(uuid: String) = {
         println("Checking health endpoint(s) for existence of consumer uuid")
         // get /health endpoint(s) and ensure it contains the new uuid
         val healthUrls: Array[String] = System.getProperty("health_url").split("\\s*,\\s*").filterNot(_.isEmpty)
         assert(healthUrls.size != 0)

         retry({
             val uuids: Array[(String, JsValue)] = healthUrls.flatMap(u => {
                 val response = RestAssured.given().config(sslconfig).get(u)
                 assert(response.statusCode() == 200)

                 response.asString()
                   .parseJson
                   .asJsObject
                   .getFields("consumers")
                   .head
                   .convertTo[JsArray]
                   .elements
                   .flatMap(c => {
                       val consumer = c.asJsObject.fields.head
                       consumer match {
                           case (u, v) if u == uuid && v.asJsObject.getFields("currentState").head == "Running".toJson => Some(consumer)
                           case _ => None
                       }
                   })
             })

             assert(uuids.nonEmpty)
         }, N = 60, waitBeforeRetry = Some(1.second))
     }

     def produceMessage(topic: String, key: String, value: String) = {
         println(s"Producing message with key: $key and value: $value")
         val producer = createProducer()
         val record = new ProducerRecord(topic, key, value)
         val future = producer.send(record)

         producer.flush()
         producer.close()

         try {
           val result = future.get(60, TimeUnit.SECONDS)

           println(s"Produced message to topic: ${result.topic()} on partition: ${result.partition()} at offset: ${result.offset()} with timestamp: ${result.timestamp()}.")
         } catch {
           case e: TimeoutException =>
             fail(s"TimeoutException received waiting for message to be produced to topic: $topic with key: $key and value: $value. ${e.getMessage}")
           case e: Exception => throw e
         }
     }
 }

 object KafkaUtils {
     val consumerInitTime = 10000 // ms

     def asKafkaProducerProps(props : Map[String,Object]) : Properties = {
         val requiredKeys = List("brokers",
                                 "user",
                                 "password",
                                 "key.serializer",
                                 "value.serializer",
                                 "security.protocol",
                                 "max.request.size")

         val propertyMap = props.filterKeys(
             requiredKeys.contains(_)
         ).map(
             tuple =>
                 tuple match {
                     // transform "brokers" key to "bootstrap.servers"
                     case (k, v) if k == "brokers" => ("bootstrap.servers", v.asInstanceOf[List[String]].mkString(","))
                     case _ => tuple
                 }
         )

         val kafkaProducerProps = new Properties()
         for ((k, v) <- propertyMap) kafkaProducerProps.put(k, v)

         kafkaProducerProps
     }

     def messagesInActivation(activation : JsObject, field: String, value: String) : Array[JsObject] = {
         val messages = JsHelpers.getFieldPath(activation, "response", "result", "messages").getOrElse(JsArray.empty).convertTo[Array[JsObject]]
         messages.filter {
             JsHelpers.getFieldPath(_, field) == Some(value.toJson)
         }
     }

     private def initializeMessageHub() = {
         // get the vcap stuff
         var credentials = TestUtils.getCredentials("message_hub")

         // initialize the set of tuples to go into the resulting Map
         val user = ("user", credentials.get("user").getAsString())
         val password = ("password", credentials.get("password").getAsString())
         val kafka_admin_url = ("kafka_admin_url", credentials.get("kafka_admin_url").getAsString())
         val api_key = ("api_key", credentials.get("api_key").getAsString())
         val security_protocol = ("security.protocol", "SASL_SSL");
         val keySerializer = ("key.serializer", "org.apache.kafka.common.serialization.StringSerializer");
         val valueSerializer = ("value.serializer", "org.apache.kafka.common.serialization.StringSerializer");
         val maxRequestSize = ("max.request.size", "3000000");
         var brokerList = new ListBuffer[String]()
         val jsonArray = credentials.get("kafka_brokers_sasl").getAsJsonArray()
         val brokerIterator = jsonArray.iterator()
         while(brokerIterator.hasNext()) {
             val current = brokerIterator.next().getAsString
             brokerList += current
         }

         val brokers = ("brokers", brokerList.toList)

         System.setProperty("java.security.auth.login.config", "")
         setMessageHubSecurityConfiguration(user._2, password._2)

         Map(user, password, kafka_admin_url, api_key, brokers, security_protocol, keySerializer, valueSerializer, maxRequestSize)
     }

     private def setMessageHubSecurityConfiguration(user: String, password: String) = {
         val map = new HashMap[String, String]()
         map.put("serviceName", "kafka")
         map.put("username", user)
         map.put("password", password)
         Configuration.setConfiguration(new Configuration()
         {
             def getAppConfigurationEntry(name: String): Array[AppConfigurationEntry] = Array(
     	          new AppConfigurationEntry (
     	              "com.ibm.messagehub.login.MessageHubLoginModule",
      			          AppConfigurationEntry.LoginModuleControlFlag.REQUIRED, map))
         })
     }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You 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 system.utils

	import java.util.HashMap
	import java.util.Properties
	import java.util.concurrent.{TimeUnit, TimeoutException}

	import io.restassured.RestAssured
	import io.restassured.config.{RestAssuredConfig, SSLConfig}
	import javax.security.auth.login.Configuration
	import javax.security.auth.login.AppConfigurationEntry
	import org.apache.kafka.clients.producer.KafkaProducer

	import scala.collection.mutable.ListBuffer
	import spray.json.DefaultJsonProtocol._
	import spray.json._
	import system.packages.ActionHelper._
	import org.apache.openwhisk.utils.JsHelpers

	import scala.concurrent.duration.DurationInt
	import common.TestHelpers
	import common.TestUtils
	import common.WskTestHelpers
	import common.ActivationResult
	import org.apache.openwhisk.utils.retry
	import org.apache.kafka.clients.producer.ProducerRecord

	trait KafkaUtils extends TestHelpers with WskTestHelpers {
	lazy val messageHubProps = KafkaUtils.initializeMessageHub()

	def createProducer() : KafkaProducer[String, String] = {
	// currently only supporting MH
	new KafkaProducer[String, String](KafkaUtils.asKafkaProducerProps(this.messageHubProps))
	}

	def apply(key : String) = {
	this.messageHubProps.getOrElse(key, "")
	}

	def getAsJson(key : String) = {
	key match {
	case key if key == "brokers" => this(key).asInstanceOf[List[String]].toJson
	case key => this(key).asInstanceOf[String].toJson
	}
	}

	val sslconfig = {
	val inner = new SSLConfig().allowAllHostnames()
	val config = inner.relaxedHTTPSValidation()
	new RestAssuredConfig().sslConfig(config)
	}

	def createTrigger(assetHelper: AssetCleaner, name: String, parameters: Map[String, spray.json.JsValue]): String = {
	println(s"Creating trigger $name")

	val feedCreationResult = assetHelper.withCleaner(wsk.trigger, name) {
	(trigger, _) =>
	trigger.create(name, feed = Some(s"/whisk.system/messaging/messageHubFeed"), parameters = parameters)
	}

	val activation = wsk.parseJsonString(feedCreationResult.stdout.substring(0, feedCreationResult.stdout.indexOf("ok: created trigger"))).convertTo[ActivationResult]

	// should be successful
	activation.response.success shouldBe true

	// It takes a moment for the consumer to fully initialize.
	println("Giving the consumer a moment to get ready")
	Thread.sleep(KafkaUtils.consumerInitTime)

	val uuid = activation.response.result.get.fields.get("uuid").get.toString().replaceAll("\"", "")
	consumerExists(uuid)

	uuid
	}


	def consumerExists(uuid: String) = {
	println("Checking health endpoint(s) for existence of consumer uuid")
	// get /health endpoint(s) and ensure it contains the new uuid
	val healthUrls: Array[String] = System.getProperty("health_url").split("\\s,\\s").filterNot(_.isEmpty)
	assert(healthUrls.size != 0)

	retry({
	val uuids: Array[(String, JsValue)] = healthUrls.flatMap(u => {
	val response = RestAssured.given().config(sslconfig).get(u)
	assert(response.statusCode() == 200)

	response.asString()
	.parseJson
	.asJsObject
	.getFields("consumers")
	.head
	.convertTo[JsArray]
	.elements
	.flatMap(c => {
	val consumer = c.asJsObject.fields.head
	consumer match {
	case (u, v) if u == uuid && v.asJsObject.getFields("currentState").head == "Running".toJson => Some(consumer)
	case _ => None
	}
	})
	})

	assert(uuids.nonEmpty)
	}, N = 60, waitBeforeRetry = Some(1.second))
	}

	def produceMessage(topic: String, key: String, value: String) = {
	println(s"Producing message with key: $key and value: $value")
	val producer = createProducer()
	val record = new ProducerRecord(topic, key, value)
	val future = producer.send(record)

	producer.flush()
	producer.close()

	try {
	val result = future.get(60, TimeUnit.SECONDS)

	println(s"Produced message to topic: ${result.topic()} on partition: ${result.partition()} at offset: ${result.offset()} with timestamp: ${result.timestamp()}.")
	} catch {
	case e: TimeoutException =>
	fail(s"TimeoutException received waiting for message to be produced to topic: $topic with key: $key and value: $value. ${e.getMessage}")
	case e: Exception => throw e
	}
	}
	}

	object KafkaUtils {
	val consumerInitTime = 10000 // ms

	def asKafkaProducerProps(props : Map[String,Object]) : Properties = {
	val requiredKeys = List("brokers",
	"user",
	"password",
	"key.serializer",
	"value.serializer",
	"security.protocol",
	"max.request.size")

	val propertyMap = props.filterKeys(
	requiredKeys.contains(_)
	).map(
	tuple =>
	tuple match {
	// transform "brokers" key to "bootstrap.servers"
	case (k, v) if k == "brokers" => ("bootstrap.servers", v.asInstanceOf[List[String]].mkString(","))
	case _ => tuple
	}
	)

	val kafkaProducerProps = new Properties()
	for ((k, v) <- propertyMap) kafkaProducerProps.put(k, v)

	kafkaProducerProps
	}

	def messagesInActivation(activation : JsObject, field: String, value: String) : Array[JsObject] = {
	val messages = JsHelpers.getFieldPath(activation, "response", "result", "messages").getOrElse(JsArray.empty).convertTo[Array[JsObject]]
	messages.filter {
	JsHelpers.getFieldPath(_, field) == Some(value.toJson)
	}
	}

	private def initializeMessageHub() = {
	// get the vcap stuff
	var credentials = TestUtils.getCredentials("message_hub")

	// initialize the set of tuples to go into the resulting Map
	val user = ("user", credentials.get("user").getAsString())
	val password = ("password", credentials.get("password").getAsString())
	val kafka_admin_url = ("kafka_admin_url", credentials.get("kafka_admin_url").getAsString())
	val api_key = ("api_key", credentials.get("api_key").getAsString())
	val security_protocol = ("security.protocol", "SASL_SSL");
	val keySerializer = ("key.serializer", "org.apache.kafka.common.serialization.StringSerializer");
	val valueSerializer = ("value.serializer", "org.apache.kafka.common.serialization.StringSerializer");
	val maxRequestSize = ("max.request.size", "3000000");
	var brokerList = new ListBuffer[String]()
	val jsonArray = credentials.get("kafka_brokers_sasl").getAsJsonArray()
	val brokerIterator = jsonArray.iterator()
	while(brokerIterator.hasNext()) {
	val current = brokerIterator.next().getAsString
	brokerList += current
	}

	val brokers = ("brokers", brokerList.toList)

	System.setProperty("java.security.auth.login.config", "")
	setMessageHubSecurityConfiguration(user._2, password._2)

	Map(user, password, kafka_admin_url, api_key, brokers, security_protocol, keySerializer, valueSerializer, maxRequestSize)
	}

	private def setMessageHubSecurityConfiguration(user: String, password: String) = {
	val map = new HashMap[String, String]()
	map.put("serviceName", "kafka")
	map.put("username", user)
	map.put("password", password)
	Configuration.setConfiguration(new Configuration()
	{
	def getAppConfigurationEntry(name: String): Array[AppConfigurationEntry] = Array(
	new AppConfigurationEntry (
	"com.ibm.messagehub.login.MessageHubLoginModule",
	AppConfigurationEntry.LoginModuleControlFlag.REQUIRED, map))
	})
	}
	}