<commit_before>package is.hail

import is.hail.annotations.{Region, RegionPool}
import is.hail.backend.BroadcastValue
import is.hail.backend.spark.SparkBackend
import is.hail.expr.ir.ExecuteContext
import is.hail.utils.{ExecutionTimer, using}
import is.hail.io.fs.FS
import org.apache.spark.SparkContext
import org.scalatest.testng.TestNGSuite
import org.testng.ITestContext
import org.testng.annotations.{AfterMethod, BeforeClass, BeforeMethod}

object HailSuite {
  def withSparkBackend(): HailContext = {
    val backend = SparkBackend(
      sc = new SparkContext(
        SparkBackend.createSparkConf(
          appName = "Hail.TestNG",
          master = System.getProperty("hail.master"),
          local = "local[2]",
          blockSize = 0)
          .set("spark.unsafe.exceptionOnMemoryLeak", "true")),
      tmpdir = "/tmp",
      localTmpdir = "file:///tmp")
    HailContext(backend, logFile = "/tmp/hail.log")
  }

  lazy val hc: HailContext = {
    val hc = withSparkBackend()
    hc.flags.set("lower", "1")
    hc.checkRVDKeys = true
    hc
  }
}

class HailSuite extends TestNGSuite {
  def hc: HailContext = HailSuite.hc

  @BeforeClass def ensureHailContextInitialized() { hc }

  def backend: SparkBackend = hc.sparkBackend("HailSuite.backend")

  def sc: SparkContext = backend.sc

  def fs: FS = backend.fs

  def fsBc: BroadcastValue[FS] = fs.broadcast

  var timer: ExecutionTimer = _

  var ctx: ExecuteContext = _

  var pool: RegionPool = _

  @BeforeMethod
  def setupContext(context: ITestContext): Unit = {
    assert(timer == null)
    timer = new ExecutionTimer("HailSuite")
    assert(ctx == null)
    pool = RegionPool()
    ctx = new ExecuteContext(backend.tmpdir, backend.localTmpdir, backend, fs, Region(pool=pool), timer, null)
  }

  @AfterMethod
  def tearDownContext(context: ITestContext): Unit = {
    ctx.close()
    ctx = null
    timer.finish()
    timer = null
    pool.close()
  }

  def withExecuteContext[T]()(f: ExecuteContext => T): T = {
    ExecutionTimer.logTime("HailSuite.withExecuteContext") { timer =>
      hc.sparkBackend("HailSuite.withExecuteContext").withExecuteContext(timer)(f)
    }
  }
}
<commit_msg>[query] Better testng sparkcontext handling (#10268)<commit_after>package is.hail

import is.hail.annotations.{Region, RegionPool}
import is.hail.backend.BroadcastValue
import is.hail.backend.spark.SparkBackend
import is.hail.expr.ir.ExecuteContext
import is.hail.utils.{ExecutionTimer, using}
import is.hail.io.fs.FS
import org.apache.spark.SparkContext
import org.scalatest.testng.TestNGSuite
import org.testng.ITestContext
import org.testng.annotations.{AfterMethod, BeforeClass, BeforeMethod}

object HailSuite {
  def withSparkBackend(): HailContext = {
    val backend = SparkBackend(
      sc = new SparkContext(
        SparkBackend.createSparkConf(
          appName = "Hail.TestNG",
          master = System.getProperty("hail.master"),
          local = "local[2]",
          blockSize = 0)
          .set("spark.unsafe.exceptionOnMemoryLeak", "true")),
      tmpdir = "/tmp",
      localTmpdir = "file:///tmp")
    HailContext(backend, logFile = "/tmp/hail.log")
  }

  lazy val hc: HailContext = {
    val hc = withSparkBackend()
    hc.flags.set("lower", "1")
    hc.checkRVDKeys = true
    hc
  }
}

class HailSuite extends TestNGSuite {
  def hc: HailContext = HailSuite.hc

  @BeforeClass def ensureHailContextInitialized() { hc }

  def backend: SparkBackend = hc.sparkBackend("HailSuite.backend")

  def sc: SparkContext = backend.sc

  def fs: FS = backend.fs

  def fsBc: BroadcastValue[FS] = fs.broadcast

  var timer: ExecutionTimer = _

  var ctx: ExecuteContext = _

  var pool: RegionPool = _

  @BeforeMethod
  def setupContext(context: ITestContext): Unit = {
    assert(timer == null)
    timer = new ExecutionTimer("HailSuite")
    assert(ctx == null)
    pool = RegionPool()
    ctx = new ExecuteContext(backend.tmpdir, backend.localTmpdir, backend, fs, Region(pool=pool), timer, null)
  }

  @AfterMethod
  def tearDownContext(context: ITestContext): Unit = {
    ctx.close()
    ctx = null
    timer.finish()
    timer = null
    pool.close()

    if (backend.sc.isStopped)
      throw new RuntimeException(s"method stopped spark context!")
  }

  def withExecuteContext[T]()(f: ExecuteContext => T): T = {
    ExecutionTimer.logTime("HailSuite.withExecuteContext") { timer =>
      hc.sparkBackend("HailSuite.withExecuteContext").withExecuteContext(timer)(f)
    }
  }
}

{'fraction_non_alphanumeric': 0.1189358372456964, 'fraction_numerical': 0.0024591996422982337, 'mean_word_length': 3.5513733468972535, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2262612', 'n_tokens_mistral': 1507, 'n_tokens_neox': 1461, 'n_words': 307}