Mock spark.read().format().options().load() but not spark.table()

Question

I have the following function which I need to test -

private HashMap<String, Dataset<Row>> getDataSources(SparkSession spark) { HashMap<String, Dataset<Row>> ds = new HashMap<String, Dataset<Row>>(); Dataset<Row>dimTenant = spark.table(dbName + "." + SparkConstants.DIM_TENANT) .select("tenant_key", "tenant_id"); Map<String, String> options = new HashMap<>(); mockOptions.put("table", bookValueTable); mockOptions.put("zkUrl", zkUrl); Dataset<Row> bookValue = spark.read().format("org.apache.phoenix.spark") .options(options) .load(); ds.put("dimTenant", dimTenant); ds.put("bookValue", bookValue); return ds; }

In this case, I actually need to execute spark.table but need to mock the output of spark.read().format(formatParam).options(optionsParam).load() based on formatParam and optionsParam

How can I achieve this?

Initially I started with mocking DataframeReader.class with deep stub (answer) but turns out spark.table itself calls spark.read().table hence spark.table got impacted. Then I tried spying the spark.read() object but since a new object is generated during the call, hence it also didn't work.

"but turns out spark.table itself calls spark.read().table hence" – if it's mocked, it's not calling anything. A Mockito Mock is a lifeless hull that doesn't do anything, unless you tell it to. — knittl
– knittl, Commented Dec 3, 2024 at 10:56
@knittl that's correct, that's why spark.table started giving me null. but i want that spark.table should run as it should while the other spark.read.format().options().load() should return what i specify it to, can you help me out here? two things i have understood is that.. probably spying is the right approach here and we might have to use whenNew of powermockito but i'm not exactly sure how to do it? — ujjawal
– ujjawal, Commented Dec 3, 2024 at 11:11

Kashyap · Accepted Answer · 2024-12-03 16:58:18Z

You can

provide your own implementation of spark.read().table() (i.e. DataframeReader.table())
check the table-name argument, if the argument is dbName + "." + SparkConstants.DIM_TENANT then call real spark, else return dummy data/null.

See Stubbing with callbacks

Collectives™ on Stack Overflow

Mock spark.read().format().options().load() but not spark.table()

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related