Configure GroovySourceSet via plugin written in Java

I’m currently trying to configure the Java and Groovy SourceDirectorySets for the test SourceSet in a given project via a plugin in buildSrc/ written in Java. Specifically, I’m trying to add multiple source directories to the given SourceDirectorySet. To do this for the Java SourceDirectorySet is pretty straightforward since I can do something like the following:

@Override
public void apply(Project project) {
  project.getPlugins().withType(JavaBasePlugin.class, javaPlugin -> {
    JavaPluginConvention javaPluginConvention = project.getConvention().getPlugin(JavaPluginConvention.class);
    SourceSetContainer sourceSets = javaPluginConvention.getSourceSets();
    SourceSet testSourceSet = sourceSets.getByName(SourceSet.TEST_SOURCE_SET_NAME);
    SourceDirectorySet testJavaSourceDirectorySet = testSourceSet.getJava();
  }
}

Getting access to the Groovy SourceDirectorySet is a bit more difficult, however, since I have to first get access to a GroovySourceSet instance. It looks like the source code ties the GroovySourceSet to it’s corresponding SourceSet via the DslObject class, but it’s not part of the public API. Any ideas on how I might be able to get access to a GroovySourceSet instance from the Java code in my plugin? Would this ‘just work’ if I wrote my plugin in Groovy instead and called testSourceSet.groovy?