<commit_before>package com.databricks.spark.redshift

import java.io.File
import java.sql.{Connection, PreparedStatement}

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.mapreduce.InputFormat
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.jdbc.JDBCWrapper
import org.apache.spark.sql.sources._
import org.apache.spark.sql.{SaveMode, Row, SQLContext}
import org.apache.spark.{SparkConf, SparkContext}
import org.scalamock.scalatest.MockFactory
import org.scalatest.{BeforeAndAfterAll, FunSuite, Matchers}

import scala.util.matching.Regex

/**
 * Tests main DataFrame loading and writing functionality
 */
class RedshiftSourceSuite
  extends FunSuite
  with Matchers
  with MockFactory
  with BeforeAndAfterAll {

  /**
   * Temporary folder for unloading data to
   */
  val tempDir = {
    var dir = File.createTempFile("spark_redshift_tests", "")
    dir.delete()
    dir.mkdirs()
    dir.deleteOnExit()
    dir.toURI.toString
  }

  /**
   * A text file containing fake unloaded Redshift data of all supported types
   */
  val testData = new File("src/test/resources/redshift_unload_data.txt").toURI.toString

  /**
   * Expected parsed output corresponding to the output of testData.
   */
  val expectedData =
    Array(
      Row(1.toByte, true, TestUtils.toTimestamp(2015, 6, 1, 0, 0, 0), 1234152.123124981,
        1.0f, 42, 1239012341823719L, 23, "Unicode是樂趣", TestUtils.toTimestamp(2015, 6, 1, 0, 0, 0, 1)),
      Row(1.toByte, false, TestUtils.toTimestamp(2015, 6, 2, 0, 0, 0), 0.0, 0.0f, 42, 1239012341823719L, -13, "asdf",
        TestUtils.toTimestamp(2015, 6, 2, 0, 0, 0, 0)),
      Row(0.toByte, null, TestUtils.toTimestamp(2015, 6, 3, 0, 0, 0), 0.0, -1.0f, 4141214, 1239012341823719L, null, "f",
        TestUtils.toTimestamp(2015, 6, 3, 0, 0, 0)),
      Row(0.toByte, false, null, -1234152.123124981, 100000.0f, null, 1239012341823719L, 24, "___|_123", null),
      Row(List.fill(10)(null): _*))

  /**
   * SparkConfig for local testing
   */
  val testConfig = new SparkConf().setMaster("local[1]").setAppName("RedshiftSourceSuite")

  /**
   * Spark Context with hadoop file overridden to point at our local test data file for this suite,
   * no-matter what temp directory was generated and requested.
   */
  private var sc: SparkContext = _

  override def beforeAll(): Unit = {
    super.beforeAll()
    sc = new SparkContext(testConfig) {
      override def newAPIHadoopFile[K, V, F <: InputFormat[K, V]]
      (path: String, fClass: Class[F], kClass: Class[K],
       vClass: Class[V], conf: Configuration = hadoopConfiguration):
      RDD[(K, V)] = {
        super.newAPIHadoopFile[K, V, F](testData, fClass, kClass, vClass, conf)
      }
    }
  }

  override def afterAll(): Unit = {
    sc.stop()
    super.afterAll()
  }

  /**
   * Set up a mocked JDBCWrapper instance that expects a sequence of queries matching the given
   * regular expressions will be executed, and that the connection returned will be closed.
   */<commit_msg>Trying to fix shutdown hook error<commit_after>package com.databricks.spark.redshift

import java.io.File
import java.sql.{Connection, PreparedStatement}

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.mapreduce.InputFormat
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.jdbc.JDBCWrapper
import org.apache.spark.sql.sources._
import org.apache.spark.sql.{SaveMode, Row, SQLContext}
import org.apache.spark.{SparkConf, SparkContext}
import org.scalamock.scalatest.MockFactory
import org.scalatest.{BeforeAndAfterAll, FunSuite, Matchers}

import scala.util.matching.Regex

/**
 * Tests main DataFrame loading and writing functionality
 */
class RedshiftSourceSuite
  extends FunSuite
  with Matchers
  with MockFactory
  with BeforeAndAfterAll {

  /**
   * Temporary folder for unloading data to
   */
  val tempDir = {
    var dir = File.createTempFile("spark_redshift_tests", "")
    dir.delete()
    dir.mkdirs()
    dir.toURI.toString
  }

  /**
   * A text file containing fake unloaded Redshift data of all supported types
   */
  val testData = new File("src/test/resources/redshift_unload_data.txt").toURI.toString

  /**
   * Expected parsed output corresponding to the output of testData.
   */
  val expectedData =
    Array(
      Row(1.toByte, true, TestUtils.toTimestamp(2015, 6, 1, 0, 0, 0), 1234152.123124981,
        1.0f, 42, 1239012341823719L, 23, "Unicode是樂趣", TestUtils.toTimestamp(2015, 6, 1, 0, 0, 0, 1)),
      Row(1.toByte, false, TestUtils.toTimestamp(2015, 6, 2, 0, 0, 0), 0.0, 0.0f, 42, 1239012341823719L, -13, "asdf",
        TestUtils.toTimestamp(2015, 6, 2, 0, 0, 0, 0)),
      Row(0.toByte, null, TestUtils.toTimestamp(2015, 6, 3, 0, 0, 0), 0.0, -1.0f, 4141214, 1239012341823719L, null, "f",
        TestUtils.toTimestamp(2015, 6, 3, 0, 0, 0)),
      Row(0.toByte, false, null, -1234152.123124981, 100000.0f, null, 1239012341823719L, 24, "___|_123", null),
      Row(List.fill(10)(null): _*))

  /**
   * SparkConfig for local testing
   */
  val testConfig = new SparkConf().setMaster("local[1]").setAppName("RedshiftSourceSuite")

  /**
   * Spark Context with hadoop file overridden to point at our local test data file for this suite,
   * no-matter what temp directory was generated and requested.
   */
  private var sc: SparkContext = _

  override def beforeAll(): Unit = {
    super.beforeAll()
    sc = new SparkContext(testConfig) {
      override def newAPIHadoopFile[K, V, F <: InputFormat[K, V]]
      (path: String, fClass: Class[F], kClass: Class[K],
       vClass: Class[V], conf: Configuration = hadoopConfiguration):
      RDD[(K, V)] = {
        super.newAPIHadoopFile[K, V, F](testData, fClass, kClass, vClass, conf)
      }
    }
  }

  override def afterAll(): Unit = {
    sc.stop()
    val temp = new File(tempDir)
    temp.listFiles() foreach {
      case f => f.delete()
    }
    temp.delete()

    super.afterAll()
  }

  /**
   * Set up a mocked JDBCWrapper instance that expects a sequence of queries matching the given
   * regular expressions will be executed, and that the connection returned will be closed.
   */
{'fraction_non_alphanumeric': 0.12313982011447261, 'fraction_numerical': 0.06410466067048241, 'mean_word_length': 3.537091988130564, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 4, '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': '4825128', 'n_tokens_mistral': 2367, 'n_tokens_neox': 2030, 'n_words': 556}