From a java plugin : how do i configure an extra sourcedir on the idea plugin

I want to move this definition to a gradle plugin written in java, not groovy :

idea {
 module {
  sourceDirs += sourceSets.somesourceset.allJava.srcDirs
 }
}

How do i do that ? I tried this but that does not work:

JavaPluginConvention javaConvention = project.getConvention().getPlugin(JavaPluginConvention.class);
IdeaModel extension = project.getExtensions().getByType(IdeaModel.class);
SourceSetContainer sourceSets = javaConvention.getSourceSets();
  SourceSet mySourceSet = sourceSets.getByName(MY_SOURCE_SET_NAME);
extension.getModule().getSourceDirs().addAll(mySourceSet.getAllJava().getSrcDirs());

Usually this is an evaluation order problem. (Is ‘sourceSets.somesourceset’ already configured at the time you apply the plugin?) There are several techniques to defer evaluation of plugin code; see similar questions in the forum or on Stack Overflow. One solution is to use a ‘project.gradle.projectsEvaluated’ callback.