diff --git a/src/test/scala/ch/cern/sparkmeasure/StageMetricsTest.scala b/src/test/scala/ch/cern/sparkmeasure/StageMetricsTest.scala
index 930ee6b..82dd872 100644
--- a/src/test/scala/ch/cern/sparkmeasure/StageMetricsTest.scala
+++ b/src/test/scala/ch/cern/sparkmeasure/StageMetricsTest.scala
@@ -31,7 +31,7 @@ class StageMetricsTest extends FunSuite with BeforeAndAfterAll {
   }
 
   test("test StageMetrics") {
-    val stageMetrics = ch.cern.sparkmeasure.StageMetrics(spark)
+    val stageMetrics = StageMetrics(spark)
     stageMetrics.runAndMeasure {
       spark.sql("select count(*) from range(100) cross join range(100) cross join range(100)").show()
     }
diff --git a/src/test/scala/ch/cern/sparkmeasure/TaskMetricsTest.scala b/src/test/scala/ch/cern/sparkmeasure/TaskMetricsTest.scala
index c768893..6f0ca51 100644
--- a/src/test/scala/ch/cern/sparkmeasure/TaskMetricsTest.scala
+++ b/src/test/scala/ch/cern/sparkmeasure/TaskMetricsTest.scala
@@ -31,7 +31,7 @@ class TaskMetricsTest extends FunSuite with BeforeAndAfterAll {
   }
 
   test("test TaskMetrics") {
-    val taskMetrics = ch.cern.sparkmeasure.TaskMetrics(spark)
+    val taskMetrics = TaskMetrics(spark)
     taskMetrics.runAndMeasure {
       spark.sql("select count(*) from range(100) cross join range(100) cross join range(100)").show()
     }