How do you Unit Test a ForeachWriter implementation? - apache-spark

Question

I've been trying to setup some unit tests to verify the logic in a ForeachWriter custom implementation but am running into a bit of mocking / duplication trouble.
I'd like to Mock an injected dependency in the ForeachWriter, but my mocks seem to be duplicated during execution. Originally I thought the mocked dependencies weren't getting called, but during debug inspection I've found that multiple versions of them seem to exist (based on hashCode).
Here's some quick sample code of what I've been trying to do:
//Class I'd like to test
public class TestForeachSink extends ForeachWriter<String> {
#Inject
SomeDependency dep;
public TestForeachSink(SomeDependency dep) {
this.dep = dep;
}
#Override
public boolean open(long partitionId, long version) {
dep.doSomethingStartupRelatedOrThrow();
return true;
}
#Override
public void process(String value) {
dep.processSomething(value);
}
#Override
public void close(Throwable errorOrNull) {
dep.closeConnections();
}
}
//Testing Class
public class TestForeachSinkTests {
#Mock SomeDependency _dep;
TestForeachSink target;
#BeforeEach
public void init() {
_dep = mock(SomeDependency.class, withSettings().serializable());
target = new TestForeachSink(_dep);
}
#Test
pubic void shouldVerifyDependencyInteractions() {
//setup stream, add data to it
stream.toDS().writeStream().foreach(target).start().processAllAvailable();
//VERIFY INTERACTIONS WITH MOCK HERE
}
}
The added data runs through the stream as expected but it seems like the mock I've passed in of SomeDependency is replaced during execution with a copy. I think that makes sense if the execution is running as though it were performing on a separate worker, but I'd still like to be able to test the ForeachWriter.
Is anyone else testing this part of the code? I haven't come across any other tests for ForeachSink custom implementations but direction on moving forward would be very appreciated!