chore: Add time in httpv2 tests to reduce flakiness on build VMs

microsoft · Mar 17, 2020 · c79dd12 · c79dd12
1 parent c7eed5a
commit c79dd12
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/main/scala/com/microsoft/ml/spark/io/http/Parsers.scala b/src/main/scala/com/microsoft/ml/spark/io/http/Parsers.scala
@@ -242,7 +242,9 @@ class CustomOutputParser(val uid: String) extends HTTPOutputParser with ComplexP
   }
   def setUDF[T: TypeTag](f: HTTPResponseData => T): this.type = {
     val fromRow = HTTPResponseData.makeFromRowConverter
-    setUDF(udf({ x: Row => f(fromRow(x)) }))
+    setUDF(udf({ x: Option[Row] =>
+      x.map(r =>
+        f(fromRow(r))) }))
   }
 
   override def transform(dataset: Dataset[_]): DataFrame = {

diff --git a/src/test/scala/com/microsoft/ml/spark/core/test/base/TestBase.scala b/src/test/scala/com/microsoft/ml/spark/core/test/base/TestBase.scala
@@ -42,7 +42,7 @@ trait LinuxOnly extends TestBase {
 
 trait Flaky extends TestBase {
 
-  val retyMillis: Array[Int] = Array(0, 100, 100)
+  val retyMillis: Array[Int] = Array(0, 1000, 5000)
 
   override def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit pos: Position): Unit = {
     super.test(testName, testTags: _*) {