From 15d9283473adfef986e829e911a9c37a98b39f8e Mon Sep 17 00:00:00 2001
From: Bago Amirbekian <bago@databricks.com>
Date: Wed, 9 May 2018 16:24:30 -0700
Subject: [PATCH] Run mllib small in unit tests (#141)

* Add mllib unit test to run mllib-small.yaml.

* Check results in unit tests and fail tests if failures are present.

* Rename to be MLLibSuite instead of MLLibTest.

* Better error message on failed benchmarks.
---
 .../spark/sql/perf/mllib/MLLibSuite.scala     | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 src/test/scala/com/databricks/spark/sql/perf/mllib/MLLibSuite.scala

diff --git a/src/test/scala/com/databricks/spark/sql/perf/mllib/MLLibSuite.scala b/src/test/scala/com/databricks/spark/sql/perf/mllib/MLLibSuite.scala
new file mode 100644
index 0000000..f335bb9
--- /dev/null
+++ b/src/test/scala/com/databricks/spark/sql/perf/mllib/MLLibSuite.scala
@@ -0,0 +1,23 @@
+package com.databricks.spark.sql.perf.mllib
+
+import org.scalatest.FunSuite
+
+import org.apache.spark.sql.Row
+
+class MLLibTest extends FunSuite {
+
+  test("test MlLib benchmarks with mllib-small.yaml.") {
+    val results = MLLib.run(yamlConfig = MLLib.smallConfig)
+    val failures = results.na.drop(Seq("failure"))
+    if (failures.count() > 0) {
+      failures.select("name", "failure.*").collect().foreach {
+        case Row(name: String, error: String, message: String) =>
+          println(
+            s"""There as a failure in the benchmark for $name:
+               |  $error ${message.replace("\n", "\n  ")}
+             """.stripMargin)
+      }
+      fail("Unable to run all benchmarks successfully, see console output for more info.")
+    }
+  }
+}