Add support for checking nullability of test code

Closes gh-11

Author: wilkinsona

README.md

@@ -8,7 +8,7 @@ A Gradle plugin for verifying a Spring project's nullability at compile time. Gr
 
 The plugin can then be applied to a project in the usual manner, as shown in the following example:
 
-```
+```groovy
 plugins {
 	id "io.spring.nullability" version "<<version>>"
 }
@@ -23,9 +23,31 @@ It will verify the nullability of all `JavaCompile` tasks whose name matches `co
 A `nullability` extension is added by the plugin.
 The extension can be used to customize the versions of the dependencies that are used during verification.
 
-```
+```groovy
 nullability {
 	errorProneVersion = <<custom-version>>
 	nullAwayVersion = <<custom-version>>
 }
 ```
+
+## Types of Nullability Checking
+
+The plugin supports two types of nullability checking, `main` and `tests`.
+The former is applied by default to tasks whose name matches `compile(\\d+)?Java`.
+The latter is opt-in.
+
+Checking using `tests` differs from `main`. It adds support for AssertJ's custom `Contract` annotation and enables NullAway's `HandleTestAssertionLibraries` option.
+
+## Configuring Nullability Checking
+
+The plugin adds a `nullability` extension to the `options` of all `JavaCompile` tasks.
+The extension provides a single property, `checking`, that can be used to configure the type of null checking that is performed.
+It defaults to `main` for tasks whose name matches `compile(\\d+)?Java`.
+For all other `JavaCompile` tasks it defaults to `disabled`.
+
+The following enables `tests` nullability checking for `compileTestJava`:
+
+```groovy
+tasks.named("compileTestJava") {
+	options.nullability.checking = "tests"
+}

src/main/java/io/spring/gradle/nullability/NullabilityOptions.java

@@ -16,15 +16,19 @@
 
 package io.spring.gradle.nullability;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 
 import javax.inject.Inject;
 
 import net.ltgt.gradle.errorprone.CheckSeverity;
 import net.ltgt.gradle.errorprone.ErrorProneOptions;
 import org.gradle.api.provider.Property;
+import org.gradle.api.provider.Provider;
 import org.gradle.api.tasks.compile.JavaCompile;
 
 /**
@@ -41,41 +45,61 @@ public abstract class NullabilityOptions {
 	 */
 	@Inject
 	public NullabilityOptions(ErrorProneOptions errorProne) {
-		errorProne.getEnabled().set(getChecking().map((checking) -> checking != Checking.DISABLED));
-		errorProne.getDisableAllChecks().set(getChecking().map((checking) -> checking != Checking.DISABLED));
-		errorProne.getCheckOptions().putAll(getChecking().map((checking) -> {
-			Map<String, String> options = new LinkedHashMap<>();
-			if (checking == Checking.MAIN) {
-				options.put("NullAway:OnlyNullMarked", "true");
-				options.put("NullAway:CustomContractAnnotations", "org.springframework.lang.Contract");
-				options.put("NullAway:JSpecifyMode", "true");
-			}
-			return options;
-		}));
-		errorProne.getChecks().putAll(getChecking().map((checking) -> {
-			if (checking == Checking.MAIN) {
-				return Map.of("NullAway", CheckSeverity.ERROR);
-			}
+		Provider<Checking> checkingAsEnum = getChecking()
+			.map((string) -> Checking.valueOf(string.toUpperCase(Locale.ROOT)));
+		errorProne.getEnabled().set(checkingAsEnum.map((checking) -> checking != Checking.DISABLED));
+		errorProne.getDisableAllChecks().set(checkingAsEnum.map((checking) -> checking != Checking.DISABLED));
+		errorProne.getCheckOptions().putAll(checkingAsEnum.map(this::checkOptions));
+		errorProne.getChecks().putAll(checkingAsEnum.map(this::checks));
+	}
+
+	private Map<String, String> checkOptions(Checking checking) {
+		if (checking == Checking.DISABLED) {
 			return Collections.emptyMap();
-		}));
+		}
+		Map<String, String> options = new LinkedHashMap<>();
+		options.put("NullAway:OnlyNullMarked", "true");
+		List<String> customContractAnnotations = new ArrayList<>();
+		customContractAnnotations.add("org.springframework.lang.Contract");
+		if (checking == Checking.TESTS) {
+			customContractAnnotations.add("org.assertj.core.internal.annotation.Contract");
+		}
+		options.put("NullAway:CustomContractAnnotations", String.join(",", customContractAnnotations));
+		options.put("NullAway:JSpecifyMode", "true");
+		if (checking == Checking.TESTS) {
+			options.put("NullAway:HandleTestAssertionLibraries", "true");
+		}
+		return options;
+	}
+
+	private Map<String, CheckSeverity> checks(Checking checking) {
+		if (checking != Checking.DISABLED) {
+			return Map.of("NullAway", CheckSeverity.ERROR);
+		}
+		return Collections.emptyMap();
 	}
 
 	/**
 	 * Returns the type of checking to perform.
 	 * @return the type of checking
 	 */
-	public abstract Property<Checking> getChecking();
+	public abstract Property<String> getChecking();
 
 	/**
-	 * The type of checking to perform for the {@link JavaCompile} task.
+	 * The type of null checking to perform for the {@link JavaCompile} task.
 	 */
-	public enum Checking {
+	enum Checking {
 
 		/**
 		 * Main code nullability checking is performed.
 		 */
 		MAIN,
 
+		/**
+		 * Test code nullability checking is performed.
+		 */
+		TESTS,
+
 		/**
 		 * Nullability checking is disabled.
 		 */

src/main/java/io/spring/gradle/nullability/NullabilityPlugin.java

@@ -64,7 +64,8 @@ private void configureJavaCompilation(Project project) {
 				.getByType(ErrorProneOptions.class);
 			NullabilityOptions nullabilityOptions = ((ExtensionAware) javaCompile.getOptions()).getExtensions()
 				.create("nullability", NullabilityOptions.class, errorProneOptions);
-			nullabilityOptions.getChecking().set(compilesMainSources(javaCompile) ? Checking.MAIN : Checking.DISABLED);
+			nullabilityOptions.getChecking()
+				.set(compilesMainSources(javaCompile) ? Checking.MAIN.name() : Checking.DISABLED.name());
 		});
 	}
 

src/test/java/io/spring/gradle/nullability/NullabilityPluginIntegrationTests.java

@@ -74,11 +74,20 @@ void configuresErrorProneOnCompileJava() {
 	}
 
 	@Test
-	void disablesErrorProneOnCompileTestJava() {
+	void disablesErrorProneOnCompileTestJavaByDefault() {
 		BuildResult result = this.gradleBuild.build("checkCompileTestJava");
 		assertThat(result.getOutput()).contains("-XepCompilingTestOnlyCode");
 	}
 
+	@Test
+	void configuresErrorProneOnCompileTestJavaWhenEnabled() {
+		BuildResult result = this.gradleBuild.build("checkCompileTestJava");
+		assertThat(result.getOutput()).contains(
+				"-XepDisableAllChecks -XepCompilingTestOnlyCode -Xep:NullAway:ERROR -XepOpt:NullAway:OnlyNullMarked=true "
+						+ "-XepOpt:NullAway:CustomContractAnnotations=org.springframework.lang.Contract,org.assertj.core.internal.annotation.Contract "
+						+ "-XepOpt:NullAway:JSpecifyMode=true -XepOpt:NullAway:HandleTestAssertionLibraries=true");
+	}
+
 	@Test
 	void compileFailsForNullabilityViolationInMainCode() throws IOException {
 		writeSource("main");
@@ -93,6 +102,13 @@ void compileSucceedsForNullabilityViolationInTestCode() {
 		assertThat(result.task(":compileTestJava").getOutcome()).isEqualTo(TaskOutcome.SUCCESS);
 	}
 
+	@Test
+	void compileFailsForNullabilityViolationInTestCodeWhenCheckingIsEnabled() throws IOException {
+		writeSource("test");
+		BuildResult result = this.gradleBuild.prepareRunner("compileTestJava").buildAndFail();
+		assertThat(result.getOutput()).contains("[NullAway] assigning @Nullable expression to @NonNull field");
+	}
+
 	private void writeSource(String sourceSetName) {
 		Path projectDir = this.gradleBuild.getProjectDir().toPath();
 		Path pkg = projectDir.resolve("src/%s/java/com/example".formatted(sourceSetName));

src/test/resources/io/spring/gradle/nullability/NullabilityPluginIntegrationTests-compileFailsForNullabilityViolationInTestCodeWhenCheckingIsEnabled.gradle

@@ -0,0 +1,16 @@
+plugins {
+	id "io.spring.nullability"
+	id "java"
+}
+
+repositories {
+	mavenCentral()
+}
+
+dependencies {
+	testCompileOnly("org.jspecify:jspecify:1.0.0")
+}
+
+tasks.named("compileTestJava") {
+	options.nullability.checking = "tests"
+}

src/test/resources/io/spring/gradle/nullability/NullabilityPluginIntegrationTests-configuresErrorProneOnCompileTestJavaWhenEnabled.gradle

@@ -0,0 +1,26 @@
+import io.spring.gradle.nullability.NullabilityOptions.Checking
+
+plugins {
+	id "io.spring.nullability"
+	id "java"
+}
+
+repositories {
+	mavenCentral()
+}
+
+dependencies {
+	compileOnly("org.jspecify:jspecify:1.0.0")
+
+	testCompileOnly("org.jspecify:jspecify:1.0.0")
+}
+
+tasks.register("checkCompileTestJava") {
+	doFirst {
+		println compileTestJava.options.extensions.errorprone
+	}
+}
+
+tasks.named("compileTestJava") {
+	options.nullability.checking = "tests"
+}