模拟SparkSession进行单元测试

时间:2018-03-26 04:27:37

标签: scala unit-testing apache-spark mocking scalamock

我的spark应用程序中有一个方法可以加载MySQL数据库中的数据。该方法看起来像这样。

trait DataManager {

val session: SparkSession

def loadFromDatabase(input: Input): DataFrame = {
            session.read.jdbc(input.jdbcUrl, s"(${input.selectQuery}) T0",
              input.columnName, 0L, input.maxId, input.parallelism, input.connectionProperties)
    }
}

除了执行jdbc方法并从数据库加载数据外,该方法不做任何其他操作。我该如何测试这种方法?标准方法是创建对象session的模拟,该对象是SparkSession的实例。但由于SparkSession有一个私有构造函数,我无法使用ScalaMock来模拟它。

这里的主要问题是我的函数是一个纯粹的副作用函数(副作用是从关系数据库中提取数据)以及我如何单独测试这个函数,因为我有模拟SparkSession的问题。

那么有什么方法可以模仿SparkSession或者其他更好的方式而不是嘲笑来测试这种方法?

2 个答案:

答案 0 :(得分:1)

在你的情况下,我建议不要嘲笑SparkSession。这或多或少会模仿整个函数(无论如何你都可以做)。如果你想测试这个函数我的建议是运行一个嵌入式数据库(如H2)并使用一个真正的SparkSession。为此,您需要向DataManager提供SparkSession。

未经测试的草图:

您的代码:

class DataManager (session: SparkSession) {
         def loadFromDatabase(input: Input): DataFrame = {
            session.read.jdbc(input.jdbcUrl, s"(${input.selectQuery}) T0",
            input.columnName, 0L, input.maxId, input.parallelism, input.connectionProperties)
         }
    }

您的测试用例:

class DataManagerTest extends FunSuite with BeforeAndAfter {
  override def beforeAll() {
    Connection conn = DriverManager.getConnection("jdbc:h2:~/test", "sa", "");
    // your insert statements goes here
    conn.close()
  }

  test ("should load data from database") {
    val dm = DataManager(SparkSession.builder().getOrCreate())
    val input = Input(jdbcUrl = "jdbc:h2:~/test", selectQuery="SELECT whateveryounedd FROM whereeveryouputit ")
    val expectedData = dm.loadFromDatabase(input)
    assert(//expectedData)
  }
}

答案 1 :(得分:0)

您可以使用Mockito Scala模拟SparkSession,如this article所示。