How to mock methods of Spark SqlContext? - scala

Question

I am writing one Spark application using Scala and I am trying to write unit tests for a method which will load data from Hive table, do some processing on it and return the result as a data frame.
Method looks as shown below:
private def filterData(SqlContext context, tableName): DataFrame = {
val table = context.table(tableName)
val selectColumnList = Seq("colA", "colB")
table.select(selectColumnList).filter(table.col("colC") > 100)
}
I would like to know how can I mock SqlContext.table() method so that I can supply some test data whenever it is called or is there any other way to achieve it using Scala?

Answer 1

Don't mock what you don't own.
When you do that, you're assuming you know how that code will behave, and therefore you can provide the result of invoking that code in your test. This assumption is likely to blow up in your face, especially when you upgrade the library version - tests pass, production breaks.
Instead, write an Adapter for it, and then use a mocked instance of it when testing units that use it. The adapter separates your code from the outside world. To test the adapter itself, you'll have to write an Integration Test, that spins up spark (or whatever implementation of the adapter) and checks that the adapter works correctly.
So, your adapter could contain the function you described above, you'd need to write an Integration test that checks it against real Spark. When you use the adapter, though, you can mock it.
trait DataProcessor {
def filterData(SqlContext context, tableName): DataFrame
}
class SparkDataProcessor extends DataProcessor {
override def filterData(SqlContext context, tableName): DataFrame = {
...
}
}
And the test for the class that uses it:
class MyThingieTest extends Spec {
"should use the data processor" >> {
val mockDataProcessor = mock[DataProcessor]
mockDataProcessor.filterData(context, tableName) returns ...
MyThingie(mockDataProcessor).doSomething must beEqualTo(...)
}
}
This way you can specify what the adapter returns.
Note - make sure to not leak 3rd party implementation in the adapter API. It should only return your data structures.
Here is another great article that talks about this very subject.