test: add scalatest assertions and fix flakiness

Author: ali-ince

common/pom.xml

@@ -46,6 +46,16 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.scalatest</groupId>
+            <artifactId>scalatest_${scala.binary.version}</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.scalatestplus</groupId>
+            <artifactId>junit-4-13_${scala.binary.version}</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

pom.xml

@@ -113,6 +113,16 @@
                 <artifactId>scala-reflect</artifactId>
                 <version>${scala.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.scalatest</groupId>
+                <artifactId>scalatest_${scala.binary.version}</artifactId>
+                <version>3.2.19</version>
+            </dependency>
+            <dependency>
+                <groupId>org.scalatestplus</groupId>
+                <artifactId>junit-4-13_${scala.binary.version}</artifactId>
+                <version>3.2.19.0</version>
+            </dependency>
             <dependency>
                 <groupId>org.slf4j</groupId>
                 <artifactId>slf4j-api</artifactId>

spark-3/pom.xml

@@ -55,6 +55,11 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.scalatest</groupId>
+            <artifactId>scalatest_${scala.binary.version}</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

spark-3/src/test/scala/org/neo4j/spark/DataSourceWriterTSE.scala

@@ -37,10 +37,13 @@ import org.neo4j.driver.summary.ResultSummary
 import org.neo4j.driver.types.IsoDuration
 import org.neo4j.driver.types.Type
 import org.neo4j.spark.util.Neo4jOptions
+import org.scalatest.matchers.must.Matchers.be
+import org.scalatest.matchers.must.Matchers.include
+import org.scalatest.matchers.must.Matchers.the
+import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
 
 import java.time.LocalTime
 import java.time.OffsetTime
-import java.time.ZoneOffset
 
 import scala.collection.JavaConverters._
 import scala.collection.immutable.ListMap
@@ -572,7 +575,7 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
     assertEquals(total, records.size)
   }
 
-  @Test(expected = classOf[SparkException])
+  @Test
   def `should throw an error because the node already exists`(): Unit = {
     SparkConnectorScalaSuiteIT.session()
       .writeTransaction(new TransactionWork[Result] {
@@ -588,24 +591,23 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
     val ds = Seq(SimplePerson("Andrea", "Santurbano")).toDS()
 
     try {
-      ds.write
-        .format(classOf[DataSource].getName)
-        .mode(SaveMode.Append)
-        .option("url", SparkConnectorScalaSuiteIT.server.getBoltUrl)
-        .option("labels", "Person")
-        .save() // we need the action to be able to trigger the exception because of the changes in Spark 3
-    } catch {
-      case sparkException: SparkException => {
-        val clientException = ExceptionUtils.getRootCause(sparkException)
-        assertNotNull(clientException)
-        assertTrue(clientException.isInstanceOf[ClientException])
-        assertTrue(
-          clientException.asInstanceOf[ClientException].code() == "Neo.ClientError.Schema.ConstraintValidationFailed"
-        )
-        throw sparkException
+      val thrown = the[SparkException] thrownBy {
+        ds.write
+          .format(classOf[DataSource].getName)
+          .mode(SaveMode.Append)
+          .option("url", SparkConnectorScalaSuiteIT.server.getBoltUrl)
+          .option("labels", "Person")
+          .save() // we need the action to be able to trigger the exception because of the changes in Spark 3
+      }
+
+      thrown.getMessage should include("org.neo4j.driver.exceptions.ClientException")
+      val rootCause = ExceptionUtils.getRootCause(thrown)
+      // root cause is not always returned as a ClientException so we pass it through pattern matching to remove flakiness
+      rootCause match {
+        case c: ClientException =>
+          c.code() should be("Neo.ClientError.Schema.ConstraintValidationFailed")
+        case _ =>
       }
-      case e: Throwable =>
-        fail(s"should be thrown a ${classOf[SparkException].getName} but is ${e.getClass.getSimpleName}")
     } finally {
       SparkConnectorScalaSuiteIT.session()
         .writeTransaction(new TransactionWork[Result] {

test-support/pom.xml

@@ -47,6 +47,14 @@
             <artifactId>powermock-module-junit4</artifactId>
             <version>${powermock.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.scalatest</groupId>
+            <artifactId>scalatest_${scala.binary.version}</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.scalatestplus</groupId>
+            <artifactId>junit-4-13_${scala.binary.version}</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.slf4j</groupId>
             <artifactId>slf4j-simple</artifactId>

test-support/src/main/scala/org/neo4j/spark/SparkConnectorScalaBaseTSE.scala

@@ -25,6 +25,7 @@ import org.neo4j.driver.Transaction
 import org.neo4j.driver.TransactionWork
 import org.neo4j.driver.summary.ResultSummary
 import org.neo4j.spark
+import org.scalatestplus.junit.JUnitSuite
 
 import java.util.concurrent.TimeUnit
 
@@ -52,7 +53,7 @@ object SparkConnectorScalaBaseTSE {
 
 }
 
-class SparkConnectorScalaBaseTSE {
+class SparkConnectorScalaBaseTSE extends JUnitSuite {
 
   val conf: SparkConf = SparkConnectorScalaSuiteIT.conf
   val ss: SparkSession = SparkConnectorScalaSuiteIT.ss