client-http/src/test/scala/org/apache/livy/client/http/HttpClientSpec.scala - incubator-livy - 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 org.apache.livy.client.http

 import java.io.{File, InputStream}
 import java.net.{InetAddress, URI}
 import java.nio.file.{Files, Paths}
 import java.util.concurrent.{Future => JFuture, _}
 import java.util.concurrent.atomic.AtomicLong
 import javax.servlet.ServletContext
 import javax.servlet.http.HttpServletRequest

 import scala.concurrent.{ExecutionContext, Future}

 import org.mockito.ArgumentCaptor
 import org.mockito.Matchers.{eq => meq, _}
 import org.mockito.Mockito._
 import org.scalatest.{BeforeAndAfterAll, FunSpecLike}
 import org.scalatra.LifeCycle
 import org.scalatra.servlet.ScalatraListener

 import org.apache.livy._
 import org.apache.livy.client.common.{BufferUtils, Serializer}
 import org.apache.livy.client.common.HttpMessages._
 import org.apache.livy.server.{AccessManager, WebServer}
 import org.apache.livy.server.interactive.{InteractiveSession, InteractiveSessionServlet}
 import org.apache.livy.server.recovery.SessionStore
 import org.apache.livy.sessions.{InteractiveSessionManager, SessionState, Spark}
 import org.apache.livy.test.jobs.Echo
 import org.apache.livy.utils.AppInfo

 /**
  * The test for the HTTP client is written in Scala so we can reuse the code in the livy-server
  * module, which implements the client session backend. The client servlet has some functionality
  * overridden to avoid creating sub-processes for each seession.
  */
 class HttpClientSpec extends FunSpecLike with BeforeAndAfterAll with LivyBaseUnitTestSuite {

   import HttpClientSpec._

   private val TIMEOUT_S = 10
   private val ID_GENERATOR = new AtomicLong()
   private val serializer = new Serializer()

   private var server: WebServer = _
   private var client: LivyClient = _

   override def beforeAll(): Unit = {
     super.beforeAll()
     server = new WebServer(new LivyConf(), "0.0.0.0", 0)

     server.context.setResourceBase("src/main/org/apache/livy/server")
     server.context.setInitParameter(ScalatraListener.LifeCycleKey,
       classOf[HttpClientTestBootstrap].getCanonicalName)
     server.context.addEventListener(new ScalatraListener)

     server.start()
   }

   override def afterAll(): Unit = {
     super.afterAll()
     if (server != null) {
       server.stop()
       server = null
     }
     if (client != null) {
       client.stop(true)
       client = null
     }
     session = null
   }

   describe("HTTP client library") {

     it("should create clients") {
       // WebServer does this internally instead of respecting "0.0.0.0", so try to use the same
       // address.
       val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}/"
       client = new LivyClientBuilder(false).setURI(new URI(uri)).build()
     }

     withClient("should run and monitor asynchronous jobs") {
       testJob(false)
     }

     withClient("should propagate errors from jobs") {
       val errorMessage = "This job throws an error."
       val (jobId, handle) = runJob(false, { id => Seq(
           new JobStatus(id, JobHandle.State.FAILED, null, errorMessage))
         })

       val error = intercept[ExecutionException] {
         handle.get(TIMEOUT_S, TimeUnit.SECONDS)
       }
       assert(error.getCause() != null)
       assert(error.getCause().getMessage().indexOf(errorMessage) >= 0)
       verify(session, times(1)).jobStatus(meq(jobId))
     }

     withClient("should run and monitor synchronous jobs") {
       testJob(false)
     }

     withClient("should add files and jars") {
       val furi = new URI("hdfs:file")
       val juri = new URI("hdfs:jar")

       client.addFile(furi).get(TIMEOUT_S, TimeUnit.SECONDS)
       client.addJar(juri).get(TIMEOUT_S, TimeUnit.SECONDS)

       verify(session, times(1)).addFile(meq(furi))
       verify(session, times(1)).addJar(meq(juri))
     }

     withClient("should upload files and jars") {
       uploadAndVerify("file")
       uploadAndVerify("jar")
     }

     withClient("should cancel jobs") {
       val (jobId, handle) = runJob(false, { id => Seq(
           new JobStatus(id, JobHandle.State.STARTED, null, null),
           new JobStatus(id, JobHandle.State.CANCELLED, null, null))
         })
       handle.cancel(true)

       intercept[CancellationException] {
         handle.get(TIMEOUT_S, TimeUnit.SECONDS)
       }

       verify(session, times(1)).cancelJob(meq(jobId))
     }

     withClient("should notify listeners of job completion") {
       val (jobId, handle) = runJob(false, { id => Seq(
           new JobStatus(id, JobHandle.State.STARTED, null, null),
           new JobStatus(id, JobHandle.State.SUCCEEDED, serialize(id), null))
         })

       val listener = mock(classOf[JobHandle.Listener[Long]])
       handle.asInstanceOf[JobHandle[Long]].addListener(listener)

       assert(handle.get(TIMEOUT_S, TimeUnit.SECONDS) === jobId)
       verify(listener, times(1)).onJobSucceeded(any(), any())
     }

     withClient("should time out handle get() call") {
       // JobHandleImpl does exponential backoff checking the result of a job. Given an initial
       // wait of 100ms, 4 iterations should result in a wait of 800ms, so the handle should at that
       // point timeout a wait of 100ms.
       val (jobId, handle) = runJob(false, { id => Seq(
           new JobStatus(id, JobHandle.State.STARTED, null, null),
           new JobStatus(id, JobHandle.State.STARTED, null, null),
           new JobStatus(id, JobHandle.State.STARTED, null, null),
           new JobStatus(id, JobHandle.State.SUCCEEDED, serialize(id), null))
         })

       intercept[TimeoutException] {
         handle.get(100, TimeUnit.MILLISECONDS)
       }

       assert(handle.get(TIMEOUT_S, TimeUnit.SECONDS) === jobId)
     }

     withClient("should handle null responses") {
       testJob(false, response = Some(null))
     }

     withClient("should connect to existing sessions") {
       var sid = client.asInstanceOf[HttpClient].getSessionId()
       val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}" +
         s"${LivyConnection.SESSIONS_URI}/$sid"
       val newClient = new LivyClientBuilder(false).setURI(new URI(uri)).build()
       newClient.stop(false)
       verify(session, never()).stop()
     }

     withClient("should tear down clients") {
       client.stop(true)
       verify(session, times(1)).stop()
     }

   }

   private def uploadAndVerify(cmd: String): Unit = {
     val f = File.createTempFile("uploadTestFile", cmd)
     val expectedStr = "Test data"
     val expectedData = expectedStr.getBytes()
     Files.write(Paths.get(f.getAbsolutePath), expectedData)
     val b = new Array[Byte](expectedData.length)
     val captor = ArgumentCaptor.forClass(classOf[InputStream])
     if (cmd == "file") {
       client.uploadFile(f).get(TIMEOUT_S, TimeUnit.SECONDS)
       verify(session, times(1)).addFile(captor.capture(), meq(f.getName))
     } else {
       client.uploadJar(f).get(TIMEOUT_S, TimeUnit.SECONDS)
       verify(session, times(1)).addJar(captor.capture(), meq(f.getName))
     }
     captor.getValue.read(b)
     assert(expectedStr === new String(b))
   }

   private def runJob(sync: Boolean, genStatusFn: Long => Seq[JobStatus]): (Long, JFuture[Int]) = {
     val jobId = java.lang.Long.valueOf(ID_GENERATOR.incrementAndGet())
     when(session.submitJob(any(classOf[Array[Byte]]), anyString())).thenReturn(jobId)

     val statuses = genStatusFn(jobId)
     val first = statuses.head
     val remaining = statuses.drop(1)
     when(session.jobStatus(meq(jobId))).thenReturn(first, remaining: _*)

     val job = new Echo(42)
     val handle = if (sync) client.run(job) else client.submit(job)
     (jobId, handle)
   }

   private def testJob(sync: Boolean, response: Option[Any] = None): Unit = {
     val (jobId, handle) = runJob(sync, { id => Seq(
         new JobStatus(id, JobHandle.State.STARTED, null, null),
         new JobStatus(id, JobHandle.State.SUCCEEDED, serialize(response.getOrElse(id)), null))
       })
     assert(handle.get(TIMEOUT_S, TimeUnit.SECONDS) === response.getOrElse(jobId))
     verify(session, times(2)).jobStatus(meq(jobId))
   }

   private def withClient(desc: String)(fn: => Unit): Unit = {
     it(desc) {
       assume(client != null, "No active client.")
       fn
     }
   }

   def serialize(value: Any): Array[Byte] = {
     BufferUtils.toByteArray(serializer.serialize(value))
   }

 }

 private object HttpClientSpec {

   // Hack warning: keep the session object available so that individual tests can mock
   // the desired behavior before making requests to the server.
   var session: InteractiveSession = _

 }

 private class HttpClientTestBootstrap extends LifeCycle {

   private implicit def executor: ExecutionContext = ExecutionContext.global

   override def init(context: ServletContext): Unit = {
     val conf = new LivyConf()
     val stateStore = mock(classOf[SessionStore])
     val sessionManager = new InteractiveSessionManager(conf, stateStore, Some(Seq.empty))
     val accessManager = new AccessManager(conf)
     val servlet = new InteractiveSessionServlet(sessionManager, stateStore, conf, accessManager) {
       override protected def createSession(req: HttpServletRequest): InteractiveSession = {
         val session = mock(classOf[InteractiveSession])
         val id = sessionManager.nextId()
         when(session.id).thenReturn(id)
         when(session.name).thenReturn(None)
         when(session.appId).thenReturn(None)
         when(session.appInfo).thenReturn(AppInfo())
         when(session.state).thenReturn(SessionState.Idle)
         when(session.proxyUser).thenReturn(None)
         when(session.kind).thenReturn(Spark)
         when(session.driverMemory).thenReturn(None)
         when(session.driverCores).thenReturn(None)
         when(session.executorMemory).thenReturn(None)
         when(session.executorCores).thenReturn(None)
         when(session.numExecutors).thenReturn(None)
         when(session.proxyUser).thenReturn(None)
         when(session.queue).thenReturn(None)
         when(session.conf).thenReturn(Map("" -> ""))
         when(session.archives).thenReturn(List())
         when(session.files).thenReturn(List())
         when(session.jars).thenReturn(List())
         when(session.pyFiles).thenReturn(List())
         when(session.stop()).thenReturn(Future.successful(()))
         when(session.ttl).thenReturn(None)
         when(session.idleTimeout).thenReturn(None)
         require(HttpClientSpec.session == null, "Session already created?")
         HttpClientSpec.session = session
         session
       }
     }

     context.mount(servlet, s"${LivyConnection.SESSIONS_URI}/*")
   }

 }
	/*
	* 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 org.apache.livy.client.http

	import java.io.{File, InputStream}
	import java.net.{InetAddress, URI}
	import java.nio.file.{Files, Paths}
	import java.util.concurrent.{Future => JFuture, _}
	import java.util.concurrent.atomic.AtomicLong
	import javax.servlet.ServletContext
	import javax.servlet.http.HttpServletRequest

	import scala.concurrent.{ExecutionContext, Future}

	import org.mockito.ArgumentCaptor
	import org.mockito.Matchers.{eq => meq, _}
	import org.mockito.Mockito._
	import org.scalatest.{BeforeAndAfterAll, FunSpecLike}
	import org.scalatra.LifeCycle
	import org.scalatra.servlet.ScalatraListener

	import org.apache.livy._
	import org.apache.livy.client.common.{BufferUtils, Serializer}
	import org.apache.livy.client.common.HttpMessages._
	import org.apache.livy.server.{AccessManager, WebServer}
	import org.apache.livy.server.interactive.{InteractiveSession, InteractiveSessionServlet}
	import org.apache.livy.server.recovery.SessionStore
	import org.apache.livy.sessions.{InteractiveSessionManager, SessionState, Spark}
	import org.apache.livy.test.jobs.Echo
	import org.apache.livy.utils.AppInfo

	/**
	* The test for the HTTP client is written in Scala so we can reuse the code in the livy-server
	* module, which implements the client session backend. The client servlet has some functionality
	* overridden to avoid creating sub-processes for each seession.
	*/
	class HttpClientSpec extends FunSpecLike with BeforeAndAfterAll with LivyBaseUnitTestSuite {

	import HttpClientSpec._

	private val TIMEOUT_S = 10
	private val ID_GENERATOR = new AtomicLong()
	private val serializer = new Serializer()

	private var server: WebServer = _
	private var client: LivyClient = _

	override def beforeAll(): Unit = {
	super.beforeAll()
	server = new WebServer(new LivyConf(), "0.0.0.0", 0)

	server.context.setResourceBase("src/main/org/apache/livy/server")
	server.context.setInitParameter(ScalatraListener.LifeCycleKey,
	classOf[HttpClientTestBootstrap].getCanonicalName)
	server.context.addEventListener(new ScalatraListener)

	server.start()
	}

	override def afterAll(): Unit = {
	super.afterAll()
	if (server != null) {
	server.stop()
	server = null
	}
	if (client != null) {
	client.stop(true)
	client = null
	}
	session = null
	}

	describe("HTTP client library") {

	it("should create clients") {
	// WebServer does this internally instead of respecting "0.0.0.0", so try to use the same
	// address.
	val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}/"
	client = new LivyClientBuilder(false).setURI(new URI(uri)).build()
	}

	withClient("should run and monitor asynchronous jobs") {
	testJob(false)
	}

	withClient("should propagate errors from jobs") {
	val errorMessage = "This job throws an error."
	val (jobId, handle) = runJob(false, { id => Seq(
	new JobStatus(id, JobHandle.State.FAILED, null, errorMessage))
	})

	val error = intercept[ExecutionException] {
	handle.get(TIMEOUT_S, TimeUnit.SECONDS)
	}
	assert(error.getCause() != null)
	assert(error.getCause().getMessage().indexOf(errorMessage) >= 0)
	verify(session, times(1)).jobStatus(meq(jobId))
	}

	withClient("should run and monitor synchronous jobs") {
	testJob(false)
	}

	withClient("should add files and jars") {
	val furi = new URI("hdfs:file")
	val juri = new URI("hdfs:jar")

	client.addFile(furi).get(TIMEOUT_S, TimeUnit.SECONDS)
	client.addJar(juri).get(TIMEOUT_S, TimeUnit.SECONDS)

	verify(session, times(1)).addFile(meq(furi))
	verify(session, times(1)).addJar(meq(juri))
	}

	withClient("should upload files and jars") {
	uploadAndVerify("file")
	uploadAndVerify("jar")
	}

	withClient("should cancel jobs") {
	val (jobId, handle) = runJob(false, { id => Seq(
	new JobStatus(id, JobHandle.State.STARTED, null, null),
	new JobStatus(id, JobHandle.State.CANCELLED, null, null))
	})
	handle.cancel(true)

	intercept[CancellationException] {
	handle.get(TIMEOUT_S, TimeUnit.SECONDS)
	}

	verify(session, times(1)).cancelJob(meq(jobId))
	}

	withClient("should notify listeners of job completion") {
	val (jobId, handle) = runJob(false, { id => Seq(
	new JobStatus(id, JobHandle.State.STARTED, null, null),
	new JobStatus(id, JobHandle.State.SUCCEEDED, serialize(id), null))
	})

	val listener = mock(classOf[JobHandle.Listener[Long]])
	handle.asInstanceOf[JobHandle[Long]].addListener(listener)

	assert(handle.get(TIMEOUT_S, TimeUnit.SECONDS) === jobId)
	verify(listener, times(1)).onJobSucceeded(any(), any())
	}

	withClient("should time out handle get() call") {
	// JobHandleImpl does exponential backoff checking the result of a job. Given an initial
	// wait of 100ms, 4 iterations should result in a wait of 800ms, so the handle should at that
	// point timeout a wait of 100ms.
	val (jobId, handle) = runJob(false, { id => Seq(
	new JobStatus(id, JobHandle.State.STARTED, null, null),
	new JobStatus(id, JobHandle.State.STARTED, null, null),
	new JobStatus(id, JobHandle.State.STARTED, null, null),
	new JobStatus(id, JobHandle.State.SUCCEEDED, serialize(id), null))
	})

	intercept[TimeoutException] {
	handle.get(100, TimeUnit.MILLISECONDS)
	}

	assert(handle.get(TIMEOUT_S, TimeUnit.SECONDS) === jobId)
	}

	withClient("should handle null responses") {
	testJob(false, response = Some(null))
	}

	withClient("should connect to existing sessions") {
	var sid = client.asInstanceOf[HttpClient].getSessionId()
	val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}" +
	s"${LivyConnection.SESSIONS_URI}/$sid"
	val newClient = new LivyClientBuilder(false).setURI(new URI(uri)).build()
	newClient.stop(false)
	verify(session, never()).stop()
	}

	withClient("should tear down clients") {
	client.stop(true)
	verify(session, times(1)).stop()
	}

	}

	private def uploadAndVerify(cmd: String): Unit = {
	val f = File.createTempFile("uploadTestFile", cmd)
	val expectedStr = "Test data"
	val expectedData = expectedStr.getBytes()
	Files.write(Paths.get(f.getAbsolutePath), expectedData)
	val b = new Array[Byte](expectedData.length)
	val captor = ArgumentCaptor.forClass(classOf[InputStream])
	if (cmd == "file") {
	client.uploadFile(f).get(TIMEOUT_S, TimeUnit.SECONDS)
	verify(session, times(1)).addFile(captor.capture(), meq(f.getName))
	} else {
	client.uploadJar(f).get(TIMEOUT_S, TimeUnit.SECONDS)
	verify(session, times(1)).addJar(captor.capture(), meq(f.getName))
	}
	captor.getValue.read(b)
	assert(expectedStr === new String(b))
	}

	private def runJob(sync: Boolean, genStatusFn: Long => Seq[JobStatus]): (Long, JFuture[Int]) = {
	val jobId = java.lang.Long.valueOf(ID_GENERATOR.incrementAndGet())
	when(session.submitJob(any(classOf[Array[Byte]]), anyString())).thenReturn(jobId)

	val statuses = genStatusFn(jobId)
	val first = statuses.head
	val remaining = statuses.drop(1)
	when(session.jobStatus(meq(jobId))).thenReturn(first, remaining: _*)

	val job = new Echo(42)
	val handle = if (sync) client.run(job) else client.submit(job)
	(jobId, handle)
	}

	private def testJob(sync: Boolean, response: Option[Any] = None): Unit = {
	val (jobId, handle) = runJob(sync, { id => Seq(
	new JobStatus(id, JobHandle.State.STARTED, null, null),
	new JobStatus(id, JobHandle.State.SUCCEEDED, serialize(response.getOrElse(id)), null))
	})
	assert(handle.get(TIMEOUT_S, TimeUnit.SECONDS) === response.getOrElse(jobId))
	verify(session, times(2)).jobStatus(meq(jobId))
	}

	private def withClient(desc: String)(fn: => Unit): Unit = {
	it(desc) {
	assume(client != null, "No active client.")
	fn
	}
	}

	def serialize(value: Any): Array[Byte] = {
	BufferUtils.toByteArray(serializer.serialize(value))
	}

	}

	private object HttpClientSpec {

	// Hack warning: keep the session object available so that individual tests can mock
	// the desired behavior before making requests to the server.
	var session: InteractiveSession = _

	}

	private class HttpClientTestBootstrap extends LifeCycle {

	private implicit def executor: ExecutionContext = ExecutionContext.global

	override def init(context: ServletContext): Unit = {
	val conf = new LivyConf()
	val stateStore = mock(classOf[SessionStore])
	val sessionManager = new InteractiveSessionManager(conf, stateStore, Some(Seq.empty))
	val accessManager = new AccessManager(conf)
	val servlet = new InteractiveSessionServlet(sessionManager, stateStore, conf, accessManager) {
	override protected def createSession(req: HttpServletRequest): InteractiveSession = {
	val session = mock(classOf[InteractiveSession])
	val id = sessionManager.nextId()
	when(session.id).thenReturn(id)
	when(session.name).thenReturn(None)
	when(session.appId).thenReturn(None)
	when(session.appInfo).thenReturn(AppInfo())
	when(session.state).thenReturn(SessionState.Idle)
	when(session.proxyUser).thenReturn(None)
	when(session.kind).thenReturn(Spark)
	when(session.driverMemory).thenReturn(None)
	when(session.driverCores).thenReturn(None)
	when(session.executorMemory).thenReturn(None)
	when(session.executorCores).thenReturn(None)
	when(session.numExecutors).thenReturn(None)
	when(session.proxyUser).thenReturn(None)
	when(session.queue).thenReturn(None)
	when(session.conf).thenReturn(Map("" -> ""))
	when(session.archives).thenReturn(List())
	when(session.files).thenReturn(List())
	when(session.jars).thenReturn(List())
	when(session.pyFiles).thenReturn(List())
	when(session.stop()).thenReturn(Future.successful(()))
	when(session.ttl).thenReturn(None)
	when(session.idleTimeout).thenReturn(None)
	require(HttpClientSpec.session == null, "Session already created?")
	HttpClientSpec.session = session
	session
	}
	}

	context.mount(servlet, s"${LivyConnection.SESSIONS_URI}/*")
	}

	}