<commit_before>package spark.debugger

import com.google.common.io.Files

import java.io.File
import java.util.Properties

import org.apache.hadoop.io._

import org.scalatest.BeforeAndAfter
import org.scalatest.FunSuite

import scala.collection.immutable.HashSet
import scala.collection.mutable.ArrayBuffer
import scala.io.Source

import spark.RDD
import spark.HashPartitioner
import spark.SparkEnv
import spark.ShuffleDependency
import spark.SparkContext
import spark.SparkContext._

class EventLoggingSuite extends FunSuite with BeforeAndAfter {
  var savedProperties: Properties = _
  before {
    savedProperties = System.getProperties()
  }
  after {
    System.setProperties(savedProperties)
  }

  def makeSparkContext(
    eventLog: File,
    enableDebugging: Boolean = true,
    enableChecksumming: Boolean = true
  ): SparkContext = {
    System.setProperty("spark.debugger.enable", enableDebugging.toString)
    System.setProperty("spark.debugger.logPath", eventLog.getPath)
    System.setProperty("spark.debugger.checksum", enableChecksumming.toString)
    val sc = new SparkContext("local", "test")
    sc
  }

  test("restore ParallelCollection from log") {
    // Initialize event log
    val tempDir = Files.createTempDir()
    val eventLog = new File(tempDir, "eventLog")

    // Make an RDD
    val sc = makeSparkContext(eventLog)
    val nums = sc.makeRDD(1 to 4)
    // Run an operation on the RDD so the scheduler becomes aware of it
    nums.collect()
    sc.stop()

    // Read it back from the event log
    val sc2 = makeSparkContext(eventLog)
    val r = new EventLogReader(sc2, Some(eventLog.getAbsolutePath))
    assert(r.rdd(nums.id).collect.toList === (1 to 4).toList)
    sc2.stop()
  }

  test("set nextRddId and nextShuffleId after restoring") {
    // Initialize event log
    val tempDir = Files.createTempDir()
    val eventLog = new File(tempDir, "eventLog")

    // Make some RDDs
    val sc = makeSparkContext(eventLog)
    val nums = sc.makeRDD(1 to 4)
    val numsMapped = nums.map(x => (x, x + 1))<commit_msg>Improve system property saving and restoration<commit_after>package spark.debugger

import com.google.common.io.Files

import java.io.File
import java.util.Properties

import org.apache.hadoop.io._

import org.scalatest.BeforeAndAfter
import org.scalatest.FunSuite

import scala.collection.immutable
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
import scala.io.Source

import spark.RDD
import spark.HashPartitioner
import spark.SparkEnv
import spark.ShuffleDependency
import spark.SparkContext
import spark.SparkContext._

class EventLoggingSuite extends FunSuite with BeforeAndAfter {
  // Reset the system properties after running each test.
  val prevProperties: mutable.Map[String, String] = new mutable.HashMap[String, String]
  val propertiesToRestore: immutable.Set[String] = immutable.HashSet(
    "spark.debugger.enable", "spark.debugger.logPath", "spark.debugger.checksum")
  after {
    for (prop <- propertiesToRestore; value <- prevProperties.get(prop)) {
      if (value == null) {
        System.clearProperty(prop)
      } else {
        System.setProperty(prop, value)
      }
    }
  }

  def makeSparkContext(
    eventLog: File,
    enableDebugging: Boolean = true,
    enableChecksumming: Boolean = true
  ): SparkContext = {
    val newProperties = Map(
      "spark.debugger.enable" -> enableDebugging.toString,
      "spark.debugger.logPath" -> eventLog.getPath,
      "spark.debugger.checksum" -> enableChecksumming.toString)
    for ((prop, newValue) <- newProperties) {
      val oldValue = System.setProperty(prop, newValue)
      prevProperties.put(prop, oldValue)
    }
    new SparkContext("local", "test")
  }

  def setSystemProperty(property: String, value: String) {
  }

  test("restore ParallelCollection from log") {
    // Initialize event log
    val tempDir = Files.createTempDir()
    val eventLog = new File(tempDir, "eventLog")

    // Make an RDD
    val sc = makeSparkContext(eventLog)
    val nums = sc.makeRDD(1 to 4)
    // Run an operation on the RDD so the scheduler becomes aware of it
    nums.collect()
    sc.stop()

    // Read it back from the event log
    val sc2 = makeSparkContext(eventLog)
    val r = new EventLogReader(sc2, Some(eventLog.getAbsolutePath))
    assert(r.rdd(nums.id).collect.toList === (1 to 4).toList)
    sc2.stop()
  }

  test("set nextRddId and nextShuffleId after restoring") {
    // Initialize event log
    val tempDir = Files.createTempDir()
    val eventLog = new File(tempDir, "eventLog")

    // Make some RDDs
    val sc = makeSparkContext(eventLog)
    val nums = sc.makeRDD(1 to 4)
    val numsMapped = nums.map(x => (x, x + 1))
{'fraction_non_alphanumeric': 0.09970238095238096, 'fraction_numerical': 0.004251700680272109, 'mean_word_length': 3.617271835132483, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 8, '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': '7318330', 'n_tokens_mistral': 1498, 'n_tokens_neox': 1467, 'n_words': 376}