1 files changed, 25 insertions, 23 deletions
diff --git a/src/test/scala/chiselTests/AnnotationNoDedup.scala b/src/test/scala/chiselTests/AnnotationNoDedup.scala
index d9cf6110..ff0005b4 100644
--- a/src/test/scala/chiselTests/AnnotationNoDedup.scala
+++ b/src/test/scala/chiselTests/AnnotationNoDedup.scala
@@ -4,7 +4,8 @@ package chiselTests
 
 import chisel3._
 import chisel3.experimental.doNotDedup
-import firrtl.FirrtlExecutionSuccess
+import chisel3.stage.{ChiselGeneratorAnnotation, ChiselStage}
+import firrtl.stage.FirrtlCircuitAnnotation
 import org.scalatest.freespec.AnyFreeSpec
 import org.scalatest.matchers.should.Matchers
 
@@ -41,33 +42,34 @@ class UsesMuchUsedModule(addAnnos: Boolean) extends Module {
 }
 
 class AnnotationNoDedup extends AnyFreeSpec with Matchers {
+  val stage = new ChiselStage
   // scalastyle:off line.size.limit
   "Firrtl provides transform that reduces identical modules to a single instance" - {
     "Annotations can be added which will prevent this deduplication for specific modules instances" in {
-      Driver.execute(Array("-X", "low", "--target-dir", "test_run_dir"), () => new UsesMuchUsedModule(addAnnos = true)) match {
-        case ChiselExecutionSuccess(_, _, Some(firrtlResult: FirrtlExecutionSuccess)) =>
-          val lowFirrtl = firrtlResult.emitted
-
-          lowFirrtl should include ("module MuchUsedModule :")
-          lowFirrtl should include ("module MuchUsedModule_1 :")
-          lowFirrtl should include ("module MuchUsedModule_3 :")
-          lowFirrtl should not include "module MuchUsedModule_2 :"
-          lowFirrtl should not include "module MuchUsedModule_4 :"
-        case _ =>
-      }
+      val lowFirrtl = stage
+        .execute(Array("-X", "low", "--target-dir", "test_run_dir"),
+                 Seq(ChiselGeneratorAnnotation(() => new UsesMuchUsedModule(addAnnos = true))))
+        .collectFirst {
+          case FirrtlCircuitAnnotation(circuit) => circuit.serialize
+        }.getOrElse(fail)
+      lowFirrtl should include ("module MuchUsedModule :")
+      lowFirrtl should include ("module MuchUsedModule_1 :")
+      lowFirrtl should include ("module MuchUsedModule_3 :")
+      lowFirrtl should not include "module MuchUsedModule_2 :"
+      lowFirrtl should not include "module MuchUsedModule_4 :"
     }
     "Turning off these annotations dedups all the occurrences" in {
-      Driver.execute(Array("-X", "low", "--target-dir", "test_run_dir"), () => new UsesMuchUsedModule(addAnnos = false)) match {
-        case ChiselExecutionSuccess(_, _, Some(firrtlResult: FirrtlExecutionSuccess)) =>
-          val lowFirrtl = firrtlResult.emitted
-
-          lowFirrtl should include ("module MuchUsedModule :")
-          lowFirrtl should not include "module MuchUsedModule_1 :"
-          lowFirrtl should not include "module MuchUsedModule_3 :"
-          lowFirrtl should not include "module MuchUsedModule_2 :"
-          lowFirrtl should not include "module MuchUsedModule_4 :"
-        case _ =>
-      }
+      val lowFirrtl = stage
+        .execute(Array("-X", "low", "--target-dir", "test_run_dir"),
+                 Seq(ChiselGeneratorAnnotation(() => new UsesMuchUsedModule(addAnnos = false))))
+        .collectFirst {
+          case FirrtlCircuitAnnotation(circuit) => circuit.serialize
+        }.getOrElse(fail)
+      lowFirrtl should include ("module MuchUsedModule :")
+      lowFirrtl should not include "module MuchUsedModule_1 :"
+      lowFirrtl should not include "module MuchUsedModule_3 :"
+      lowFirrtl should not include "module MuchUsedModule_2 :"
+      lowFirrtl should not include "module MuchUsedModule_4 :"
     }
   }
   // scalastyle:on line.size.limit