prevent `System.exit()` during test

ExitPrevention.scala

import java.security.Permission
import org.scalatest.{BeforeAndAfter, Suite}

final class ExitPreventionException(val status: Int)
  extends SecurityException(s"System.exit($status) is called")

trait ExitPrevention extends BeforeAndAfter { self: Suite =>

  before {
    ExitPrevention.begin()
  }

  after {
    ExitPrevention.end()
  }

}

object ExitPrevention {

  private[this] object NoExitSecurityManager extends SecurityManager {

    override def checkPermission(perm: Permission): Unit =
      ()

    override def checkPermission(perm: Permission, context: Object): Unit =
      ()

    override def checkExit(status: Int): Unit =
      throw new ExitPreventionException(status)

  }

  private[this] var originalSecurityManager: SecurityManager =
    _

  private[this] var preventionCount: Int =
    0

  def begin(): Unit =
    synchronized {
      if (preventionCount == 0) {
        originalSecurityManager = System.getSecurityManager
        System.setSecurityManager(NoExitSecurityManager)
      }
      preventionCount += 1
    }

  def end(): Unit =
    synchronized {
      preventionCount -= 1
      if (preventionCount == 0) {
        System.setSecurityManager(originalSecurityManager)
      }
    }

}