Extending JUnit plugin

Hi!
I've been trying to extend JUnit plugin in order to access JUnitConfiguration (extend it or access it's methods). Compiling project I get: 

JUnitConfiguration runProfile = (JUnitConfiguration) environment.getRunProfile();
                                            ^
class file for com.intellij.execution.JavaTestConfigurationWithDiscoverySupport not found

I've added JUnit in build.gradle:

intellij {
    plugins 'JUnit'
}

Maybe there's another way required to access JUnitConfiguration?

Thanks in advance!