Configuration contains additional imports that should be added to generated classes

accurest-core/src/main/groovy/io/coderate/accurest/TestGenerator.groovy

@@ -1,6 +1,7 @@
 package io.coderate.accurest
 
 import io.coderate.accurest.builder.ClassBuilder
+import io.coderate.accurest.config.AccurestConfigProperties
 import io.coderate.accurest.config.TestFramework
 import io.coderate.accurest.config.TestMode
 import io.coderate.accurest.util.NamesUtil
@@ -16,44 +17,31 @@ import static io.coderate.accurest.builder.MethodBuilder.createTestMethod
  */
 class TestGenerator {
 
-	private final String stubsBaseDirectory
-	private final String basePackageForTests
-	private final String baseClassForTests
-	private final String ruleClassForTests
-	private final TestFramework lang
 	private final String targetDirectory
-	private final TestMode testMode
+	private final AccurestConfigProperties configProperties
+	private final String stubsBaseDirectory
 
-	TestGenerator(String stubsBaseDirectory, String basePackageForTests, String baseClassForTests,
-	              String ruleClassForTests, TestFramework testFramework, TestMode testMode, String targetDirectory) {
-		this.testMode = testMode
-		this.targetDirectory = targetDirectory
-		File stubsResource = new File(stubsBaseDirectory)
+	TestGenerator(AccurestConfigProperties accurestConfigProperties) {
+		this.configProperties = accurestConfigProperties
+		this.targetDirectory = accurestConfigProperties.generatedTestSourcesDir
+		File stubsResource = new File(accurestConfigProperties.stubsBaseDirectory)
 		if (stubsResource == null) {
-			throw new IllegalStateException("Stubs directory not found under " + stubsBaseDirectory)
+			throw new IllegalStateException("Stubs directory not found under " + accurestConfigProperties.stubsBaseDirectory)
 		}
 		this.stubsBaseDirectory = stubsResource.path
-		this.basePackageForTests = basePackageForTests
-		if (testFramework == 'Spock' && !baseClassForTests) {
-			this.baseClassForTests = 'spock.lang.Specification'
-		} else {
-			this.baseClassForTests = baseClassForTests
-		}
-		this.ruleClassForTests = ruleClassForTests
-		this.lang = testFramework
 	}
 
 	public void generate() {
 		List<File> files = new File(stubsBaseDirectory).listFiles()
 		files.grep({ File file -> file.isDirectory() && containsStubs(file) }).each {
-			def testBaseDir = Paths.get(targetDirectory, NamesUtil.packageToDirectory(basePackageForTests))
+			def testBaseDir = Paths.get(targetDirectory, NamesUtil.packageToDirectory(configProperties.basePackageForTests))
 			Files.createDirectories(testBaseDir)
 			Files.write(Paths.get(testBaseDir.toString(), NamesUtil.capitalize(it.name) + getTestClassExtension()), addClass(it).bytes)
 		}
 	}
 
 	private String getTestClassExtension() {
-		return lang == TestFramework.SPOCK ? '.groovy' : '.java'
+		return configProperties.targetFramework == TestFramework.SPOCK ? '.groovy' : '.java'
 	}
 
 	boolean containsStubs(File file) {
@@ -66,32 +54,46 @@ class TestGenerator {
 	}
 
 	private String addClass(File directory) {
-		ClassBuilder clazz = createClass(NamesUtil.capitalize(NamesUtil.afterLast(directory.path, '/')), basePackageForTests, baseClassForTests, lang)
+		ClassBuilder clazz = createClass(NamesUtil.capitalize(NamesUtil.afterLast(directory.path, '/')), configProperties)
 
-		if (testMode == TestMode.MOCKMVC) {
+		if (configProperties.imports) {
+			configProperties.imports.each {
+				clazz.addImport(it)
+			}
+		}
+
+		if (configProperties.staticImports) {
+			configProperties.staticImports.each {
+				clazz.addStaticImport(it)
+			}
+		}
+
+		if (configProperties.testMode == TestMode.MOCKMVC) {
 			clazz.addStaticImport('com.jayway.restassured.module.mockmvc.RestAssuredMockMvc.*')
 		} else {
 			clazz.addStaticImport('com.jayway.restassured.RestAssured.*')
 		}
 
-		if (lang == TestFramework.JUNIT) {
+		if (configProperties.targetFramework == TestFramework.JUNIT) {
 			clazz.addImport('org.junit.Test')
 		} else {
 			clazz.addImport('groovy.json.JsonSlurper')
 		}
 
-		if (ruleClassForTests) {
+		if (configProperties.ruleClassForTests) {
 			clazz.addImport('org.junit.Rule')
-					.addRule(ruleClassForTests)
+					.addRule(configProperties.ruleClassForTests)
 		}
 
 		directory.listFiles().each {
-			clazz.addMethod(createTestMethod(it, lang))
+			clazz.addMethod(createTestMethod(it, configProperties.targetFramework))
 		}
 		return clazz.build()
 	}
 
 	public static void main(String[] args) {
-		new TestGenerator('/home/devel/projects/codearte/accurest/accurest-core/src/main/resources/stubs', 'io.test', '', '', TestFramework.SPOCK, TestMode.MOCKMVC, "").generate()
+		AccurestConfigProperties properties = new AccurestConfigProperties(stubsBaseDirectory: '/home/devel/projects/codearte/accurest/accurest-core/src/main/resources/stubs',
+		targetFramework: TestFramework.SPOCK, testMode: TestMode.MOCKMVC, basePackageForTests: 'io.test', staticImports: ['com.pupablada.Test.*'], imports: ['org.innapypa.Test'])
+		new TestGenerator(properties).generate()
 	}
 }

accurest-core/src/main/groovy/io/coderate/accurest/builder/ClassBuilder.groovy

@@ -1,5 +1,6 @@
 package io.coderate.accurest.builder
 
+import io.coderate.accurest.config.AccurestConfigProperties
 import io.coderate.accurest.config.TestFramework
 import io.coderate.accurest.util.NamesUtil
 
@@ -36,12 +37,14 @@ class ClassBuilder {
 		}
 	}
 
-	static ClassBuilder createClass(String className, String packageName, TestFramework lang) {
-		return createClass(className, packageName, null, lang)
-	}
-
-	static ClassBuilder createClass(String className, String packageName, String baseClass, TestFramework lang) {
-		return new ClassBuilder(className, packageName, baseClass, lang)
+	static ClassBuilder createClass(String className, AccurestConfigProperties properties) {
+		String baseClassForTests
+		if (properties.targetFramework == TestFramework.SPOCK && !properties.baseClassForTests) {
+			baseClassForTests = 'spock.lang.Specification'
+		} else {
+			baseClassForTests = properties.baseClassForTests
+		}
+		return new ClassBuilder(className, properties.basePackageForTests, baseClassForTests, properties.targetFramework)
 	}
 
 	ClassBuilder addImport(String importToAdd) {

accurest-core/src/main/groovy/io/coderate/accurest/config/AccurestConfigProperties.groovy

@@ -1,17 +1,16 @@
-package io.codearte.accurest.plugin
-
-import io.coderate.accurest.config.TestFramework
-import io.coderate.accurest.config.TestMode
+package io.coderate.accurest.config
 
 /**
  * @author Jakub Kubrynski
  */
-class AccurestPluginExtension {
+class AccurestConfigProperties {
 	TestFramework targetFramework = TestFramework.SPOCK
-	TestMode testMode = TestMode.EXPLICIT
+	TestMode testMode = TestMode.MOCKMVC
 	String stubsBaseDirectory = 'mappings'
 	String basePackageForTests = 'io.codearte.accurest.tests'
 	String baseClassForTests
 	String ruleClassForTests
 	String generatedTestSourcesDir = 'build/generated-sources/accurest'
+	String[] imports = []
+	String[] staticImports = []
 }

accurest-gradle-plugin/src/main/groovy/io/codearte/accurest/plugin/AccurestGradlePlugin.groovy

@@ -1,6 +1,7 @@
 package io.codearte.accurest.plugin
 
 import io.coderate.accurest.TestGenerator
+import io.coderate.accurest.config.AccurestConfigProperties
 import org.gradle.api.Plugin
 import org.gradle.api.Project
 
@@ -15,7 +16,7 @@ class AccurestGradlePlugin implements Plugin<Project> {
 
 	@Override
 	void apply(Project project) {
-		AccurestPluginExtension extension = project.extensions.create('accurest', AccurestPluginExtension)
+		AccurestConfigProperties extension = project.extensions.create('accurest', AccurestConfigProperties)
 
 		project.compileTestGroovy.dependsOn(TASK_NAME)
 
@@ -27,11 +28,8 @@ class AccurestGradlePlugin implements Plugin<Project> {
 			}
 
 			try {
-				String resourceDirectory = project.projectDir.path + '/src/test/resources/' + extension.stubsBaseDirectory
-				TestGenerator generator = new TestGenerator(resourceDirectory,
-						extension.basePackageForTests, extension.baseClassForTests, extension.ruleClassForTests,
-						extension.targetFramework, extension.testMode,
-						extension.generatedTestSourcesDir)
+				extension.stubsBaseDirectory = project.projectDir.path + '/src/test/resources/' + extension.stubsBaseDirectory
+				TestGenerator generator = new TestGenerator(extension)
 				generator.generate()
 			} catch (IllegalStateException e) {
 				project.logger.error("Accurest Plugin: {}", e.getMessage())