模拟用于单元测试的SparkSession

Question

我的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或者其他比模拟更好的方法来测试这个方法呢？

Answer 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)
  }
}