Add firrtlTests.Utils methods

This adds methods for examining stdout/stderr and exit codes inside of
a Scala program. This are pulled directly from firrtlTests, but we
aren't currently publishing those anywhere that we can get at them.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>

1 files changed, 85 insertions, 0 deletions

diff --git a/src/test/scala/chiselTests/ChiselSpec.scala b/src/test/scala/chiselTests/ChiselSpec.scala
index 75fa68dd..d0a1f5c5 100644
--- a/src/test/scala/chiselTests/ChiselSpec.scala
+++ b/src/test/scala/chiselTests/ChiselSpec.scala
@@ -11,6 +11,8 @@ import chisel3.testers._
 import firrtl.options.OptionsException
 import firrtl.{AnnotationSeq, CommonOptions, ExecutionOptionsManager, FirrtlExecutionFailure, FirrtlExecutionSuccess, HasFirrtlOptions}
 import firrtl.util.BackendCompilationUtilities
+import java.io.ByteArrayOutputStream
+import java.security.Permission
 
 /** Common utility functions for Chisel unit tests. */
 trait ChiselRunners extends Assertions with BackendCompilationUtilities {
@@ -200,3 +202,86 @@ class ChiselPropSpec extends PropSpec with ChiselRunners with PropertyChecks wit
     j <- Gen.choose(0, (1 << w) - 1)
   } yield (w, i, j)
 }
+
+trait Utils {
+
+  /** Run some Scala thunk and return STDOUT and STDERR as strings.
+    * @param thunk some Scala code
+    * @return a tuple containing STDOUT, STDERR, and what the thunk returns
+    */
+  def grabStdOutErr[T](thunk: => T): (String, String, T) = {
+    val stdout, stderr = new ByteArrayOutputStream()
+    val ret = scala.Console.withOut(stdout) { scala.Console.withErr(stderr) { thunk } }
+    (stdout.toString, stderr.toString, ret)
+  }
+
+  /** Encodes a System.exit exit code
+    * @param status the exit code
+    */
+  private case class ExitException(status: Int) extends SecurityException(s"Found a sys.exit with code $status")
+
+  /** A security manager that converts calls to System.exit into [[ExitException]]s by explicitly disabling the ability of
+    * a thread to actually exit. For more information, see:
+    *   - https://docs.oracle.com/javase/tutorial/essential/environment/security.html
+    */
+  private class ExceptOnExit extends SecurityManager {
+    override def checkPermission(perm: Permission): Unit = {}
+    override def checkPermission(perm: Permission, context: Object): Unit = {}
+    override def checkExit(status: Int): Unit = {
+      super.checkExit(status)
+      throw ExitException(status)
+    }
+  }
+
+  /** Encodes a file that some code tries to write to
+    * @param the file name
+    */
+  private case class WriteException(file: String) extends SecurityException(s"Tried to write to file $file")
+
+  /** A security manager that converts writes to any file into [[WriteException]]s.
+    */
+  private class ExceptOnWrite extends SecurityManager {
+    override def checkPermission(perm: Permission): Unit = {}
+    override def checkPermission(perm: Permission, context: Object): Unit = {}
+    override def checkWrite(file: String): Unit = {
+      super.checkWrite(file)
+      throw WriteException(file)
+    }
+  }
+
+  /** Run some Scala code (a thunk) in an environment where all System.exit are caught and returned. This avoids a
+    * situation where a test results in something actually exiting and killing the entire test. This is necessary if you
+    * want to test a command line program, e.g., the `main` method of [[firrtl.options.Stage Stage]].
+    *
+    * NOTE: THIS WILL NOT WORK IN SITUATIONS WHERE THE THUNK IS CATCHING ALL [[Exception]]s OR [[Throwable]]s, E.G.,
+    * SCOPT. IF THIS IS HAPPENING THIS WILL NOT WORK. REPEAT THIS WILL NOT WORK.
+    * @param thunk some Scala code
+    * @return either the output of the thunk (`Right[T]`) or an exit code (`Left[Int]`)
+    */
+  def catchStatus[T](thunk: => T): Either[Int, T] = {
+    try {
+      System.setSecurityManager(new ExceptOnExit())
+      Right(thunk)
+    } catch {
+      case ExitException(a) => Left(a)
+    } finally {
+      System.setSecurityManager(null)
+    }
+  }
+
+  /** Run some Scala code (a thunk) in an environment where file writes are caught and the file that a program tries to
+    * write to is returned. This is useful if you want to test that some thunk either tries to write to a specific file
+    * or doesn't try to write at all.
+    */
+  def catchWrites[T](thunk: => T): Either[String, T] = {
+    try {
+      System.setSecurityManager(new ExceptOnWrite())
+      Right(thunk)
+    } catch {
+      case WriteException(a) => Left(a)
+    } finally {
+      System.setSecurityManager(null)
+    }
+  }
+
+}