Test file creation in the test source folder

I would like to create a plugin which generates some unit test code based on the source code. Therefore I would like to create a new test file in the test source folders. I managed create a file only in the source folders. I found a code snippet

PsiDirectory getTheFirstPsiDirectoryInTheProject(AnActionEvent e, Project project) {
    final IdeView view = e.getData(LangDataKeys.IDE_VIEW);
    if (view == null) {
        return null;
    }

    final ProjectRootManager rootManager = ProjectRootManager.getInstance(project);
    final ProjectFileIndex fileIndex = rootManager.getFileIndex();
    final Optional<PsiDirectory> sourceDirectory = Stream.of(view.getDirectories())
        .filter(directory -> {
            final VirtualFile virtualFile = directory.getVirtualFile();
            return fileIndex.isUnderSourceRootOfType(virtualFile, JavaModuleSourceRootTypes.SOURCES);
        })
        .findFirst();

    return sourceDirectory.orElse(null);
}

which retrieves a directory in the source folder but not in the test source folder. I tried to tweak this code to return with a directory in the test source folder but I didn't manage.
So how can I retrieve a directory in the test source folder when I know the AnActionEvent and the project reference whose type is com.intellij.openapi.project.Project?