Dagger SPI - Adding Custom Graph Validation (with KSP and tests!)

Dagger SPI - Adding Custom Graph Validation (with KSP and tests!)

·

10 min read

From the Dagger site:

The Dagger Service Provider Interface (SPI) is a mechanism to hook into Dagger’s annotation processor and access the same binding graph model that Dagger uses. With the SPI, you can write a plugin that

  • adds project-specific errors and warnings, e.g. Bindings for android.content.Context must have a @Qualifier or Bindings that implement DatabaseRelated must be scoped

  • generates extra Java source files

  • serializes the Dagger model to a resource file

  • builds visualizations of the Dagger model

  • and more!

The first bullet point is what we are after today.

Moving as many checks as possible to static analysis tools like Lint, Detekt, and SPI plugins is a great way to Shift Left and fix these types of issues before a PR is ever raised.

Rules

Let's start first by defining the custom behavior we want to validate with our plugins, and then we will get to writing tests and incorporating the plugins into our project.

The three plugins we will be writing in this post are:

  1. Make sure no un-qualified primitives are provided into the Dagger graph. Something like this would be flagged with the plugin:

    1.   @Module
        class PrimitiveModule {
          // should have a qualifier to provide a primitive
          @Provides
          fun provideString(): String {
            return "test"
          }
        }
      
  2. Make sure classes that extend from a single interface are added to the Dagger graph via @Binds instead of @Provides Something like this would not be allowed:

    1.   interface TestInterface {
          fun hello()
        }
      
        class TestClass : TestInterface {
          override fun hello() {
            TODO("Not yet implemented")
          }
        }
      
        @Singleton @Component(modules = [MyModule::class])
        interface MyComponent {
          fun testClass(): TestClass
        }
      
        @Module object MyModule {
          // should be added via @Binds instead
          @Provides @Singleton fun providesTestClass() = TestClass()
        }
      
  3. Lastly, we want to make sure classes are provided uniquely into the graph between @Inject constructors and @Provides :

    1.   @Singleton @Component
        interface MyComponent {
          fun dependency(): Dependency
          fun subcomponent(): MySubcomponent
        }
      
        @Subcomponent(modules = [MySubmodule::class])
        interface MySubcomponent {
          fun dependency(): Dependency
        }
      
        @Module object MySubmodule {
          @Provides fun providesDependency() = Dependency("Provides")
        }
      
        // Already added via @Provides, no need for @Inject constructor
        class Dependency @Inject constructor(dep: String)
      

Setup

To start, we will want to create a JVM module to host our new plugins. We will also need KSP and the following dependencies:

plugins {
    id("org.jetbrains.kotlin.jvm")
    id("com.google.devtools.ksp")
}

dependencies {
    val daggerVersion = "2.48"
    val autoServiceKspVersion = "1.1.0"
    val googleAutoServiceVersion = "1.1.1"
    val kctVersion = "0.4.0"

    implementation("com.google.dagger:dagger-spi:$daggerVersion")
    ksp("dev.zacsweers.autoservice:auto-service-ksp:$autoServiceKspVersion")
    compileOnly("com.google.auto.service:auto-service:$googleAutoServiceVersion")

    // for testing our SPI plugins
    testImplementation("dev.zacsweers.kctfork:core:$kctVersion")
    testImplementation("dev.zacsweers.kctfork:ksp:$kctVersion")
    testImplementation("com.google.dagger:dagger-compiler:$daggerVersion")
}

Within our newly created module, assuming it was called :dagger-spi-plugins , we will want to add our first plugin.

The following plugin makes sure all primitive types are provided with a qualifier:

import com.google.auto.service.AutoService
import dagger.spi.model.BindingGraph
import dagger.spi.model.BindingGraphPlugin
import dagger.spi.model.DiagnosticReporter
import javax.tools.Diagnostic

// list of class types we are interested in
private val PRIMITIVES =
  listOf(
    Integer::class.java.name,
    String::class.java.name,
    Boolean::class.java.name,
    Float::class.java.name,
    Double::class.java.name
  )

// Adds the plugin via a Service Loader with Zac Sweer's fork of auto-service
@AutoService(BindingGraphPlugin::class)
class PrimitiveValidator : BindingGraphPlugin {

  override fun visitGraph(bindingGraph: BindingGraph, diagnosticReporter: DiagnosticReporter) {
    bindingGraph
       // grab all bindings
      .bindings()
       // filter the types if they are contained in our primitives list and don't have a qualifier
      .filter { binding ->
        val key = binding.key()
        PRIMITIVES.contains(key.type().toString()) && !key.qualifier().isPresent
      }
      .forEach { binding ->
        // report an error for any present binding
        diagnosticReporter.reportBinding(
          Diagnostic.Kind.ERROR,
          binding,
          "Primitives should be annotated with any qualifier"
        )
      }
  }
}

Now we can add our 2nd plugin. This plugin will make sure classes that extend from a single interface are added to the Dagger graph via @Binds instead of @Provides :

import com.google.auto.service.AutoService
import com.google.devtools.ksp.closestClassDeclaration
import com.google.devtools.ksp.symbol.ClassKind
import dagger.spi.model.Binding
import dagger.spi.model.BindingGraph
import dagger.spi.model.BindingGraphPlugin
import dagger.spi.model.BindingKind
import dagger.spi.model.DaggerProcessingEnv
import dagger.spi.model.DiagnosticReporter
import javax.tools.Diagnostic

@AutoService(BindingGraphPlugin::class)
class BindsValidator : BindingGraphPlugin {

  override fun visitGraph(bindingGraph: BindingGraph, diagnosticReporter: DiagnosticReporter) {
    // grab all bindings
    bindingGraph.bindings()
      .asSequence()
      // filter for things being provided via @Provides
      .filter { it.kind() == BindingKind.PROVISION }
      .forEach { provisionedBinding ->
        // grab all classes with a single super type that is an interface
        val classBindingsWithInterfaceSuperTypes = bindingGraph
          .bindings(provisionedBinding.key())
          .asSequence()
          .filter { binding ->
            // we are only interested in the KSP pass, not javac
            if (binding.key().type().backend() == DaggerProcessingEnv.Backend.KSP) {
              binding.key().type().ksp().declaration.closestClassDeclaration()?.let { closestClassDeclaration ->
                // if the binding we are visiting is a class
                closestClassDeclaration.classKind == ClassKind.CLASS &&
                  // extending from exactly 1 thing
                  closestClassDeclaration.superTypes.count() == 1 &&
                  // and that thing is an interface
                  closestClassDeclaration.superTypes.first().resolve().declaration.closestClassDeclaration()?.classKind == ClassKind.INTERFACE
              } ?: false
            } else {
              false
            }
          }
          .toList()

        if (classBindingsWithInterfaceSuperTypes.isNotEmpty()) {
          // report an error for each instance
          classBindingsWithInterfaceSuperTypes.forEach {
            diagnosticReporter.reportBinding(
              Diagnostic.Kind.ERROR,
              it,
              "${it.key().type().ksp().declaration.simpleName.getShortName()} should be added to the graph via @Binds instead of @Provides since it extends a single interface"
            )
          }
        }
      }
  }
}

The last plugin we want to add is one that makes sure classes are provided uniquely into the graph between @Inject constructors and @Provides :

import com.google.auto.service.AutoService
import dagger.spi.model.Binding
import dagger.spi.model.BindingGraph
import dagger.spi.model.BindingGraphPlugin
import dagger.spi.model.BindingKind
import dagger.spi.model.DiagnosticReporter
import javax.tools.Diagnostic

@AutoService(BindingGraphPlugin::class)
class UniqueInjectBindingValidator : BindingGraphPlugin {

  override fun visitGraph(bindingGraph: BindingGraph, diagnosticReporter: DiagnosticReporter) {
    bindingGraph
      // grab all bindings
      .bindings()
      .asSequence()
      // filter for bindings being added via @Inject constructors
      .filter { it.kind() == BindingKind.INJECTION }
      .forEach { injectBinding: Binding ->
        // grab other bindings that are not @Inject consturcotrs or member injection
        val otherBindings =
          bindingGraph
            .bindings(injectBinding.key())
            .asSequence()
            // Filter out other @Inject bindings
            .filter { it.kind() != BindingKind.INJECTION }
            // Filter out other member injection sites
            .filter { it.kind() != BindingKind.MEMBERS_INJECTION }
            .toList()

        // we are left with a list of bindings provided via @Provides
        // report an error for each one
        if (otherBindings.isNotEmpty()) {
          diagnosticReporter.reportBinding(
            Diagnostic.Kind.ERROR,
            injectBinding,
            "@Inject constructor has duplicates other bindings: ${otherBindings.joinToString()}"
          )
        }
      }
  }
}

Phew, I'm a little tired after writing those plugins. Let's take a a small break and look at a cute puppy.

National Puppy Day 2022: 30 cute puppy photos to make you smile

Now that we are feeling refreshed, let's write some unit tests for our plugins! In the test source set for our :dagger-spi-plugins module, we will first add some utility methods to better facilitate our testing.

We will be using Zac Sweer's fork of the Kotlin Compile Testing library to test our code output, and these utility functions just make our test setup easier.

import com.tschuchort.compiletesting.JvmCompilationResult
import com.tschuchort.compiletesting.KotlinCompilation
import com.tschuchort.compiletesting.SourceFile
import com.tschuchort.compiletesting.kspArgs
import com.tschuchort.compiletesting.kspSourcesDir
import com.tschuchort.compiletesting.kspWithCompilation
import com.tschuchort.compiletesting.symbolProcessorProviders
import dagger.internal.codegen.KspComponentProcessor
import org.intellij.lang.annotations.Language
import java.io.File
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi

internal interface KotlinCompiler {
  fun compile(vararg sourceFiles: SourceFile): CompilationResult
}

@OptIn(ExperimentalCompilerApi::class)
internal class KspKotlinCompiler(tempDir: File) : KotlinCompiler {

  private val compiler =
    KotlinCompilation().apply {
      inheritClassPath = true
      symbolProcessorProviders =
        listOf(
          KspComponentProcessor.Provider.withTestPlugins(
            // our 3 custom plugins written earlier
            BindsValidator(),
            PrimitiveValidator(),
            UniqueInjectBindingValidator()
          ),
        )
      kspArgs = getKspArguments()
      kspWithCompilation = true
      verbose = false
      workingDir = tempDir
    }

  override fun compile(vararg sourceFiles: SourceFile): CompilationResult {
    compiler.sources = sourceFiles.asList()
    return CompilationResult(
      compiler.compile(),
      findGeneratedFiles(compiler.kspSourcesDir),
      compiler.workingDir.resolve("sources"),
    )
  }
}

internal data class CompilationResult
@OptIn(ExperimentalCompilerApi::class)
constructor(
  val result: JvmCompilationResult,
  val generatedFiles: List<File>,
  val sourcesDir: File,
)

private fun findGeneratedFiles(file: File): List<File> {
  return file.walkTopDown().filter { it.isFile }.toList()
}

private fun getKspArguments() =
  mutableMapOf(
    "dagger.fullBindingGraphValidation" to "ERROR",
  )

internal fun createSource(@Language("kotlin") code: String): SourceFile {
  return SourceFile.kotlin("${code.findFullQualifiedName()}.kt", code)
}

private fun String.findFullQualifiedName(): String {
  val packageRegex = "package (.*)".toRegex()
  val packageName = packageRegex.find(this)?.groupValues?.get(1)?.trim()
  val objectRegex = "(abstract )?(class|interface|object) ([^ ]*)".toRegex()
  val objectMatcher = checkNotNull(objectRegex.find(this)) { "No class/interface/object found" }
  val objectName = objectMatcher.groupValues[3]
  return if (packageName != null) {
    "${packageName.replace(".", "/")}/$objectName"
  } else {
    objectName
  }
}

Now that we have some helper functions, let's write the first test for the PrimitiveValidator plugin:

import com.tschuchort.compiletesting.KotlinCompilation
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
import org.junit.Rule
import org.junit.Test
import org.junit.rules.TemporaryFolder

@OptIn(ExperimentalCompilerApi::class)
class PrimitiveValidatorTest {

  @JvmField @Rule var folder = TemporaryFolder()

  @Test
  fun `primitives should not be provided without a qualifier`() {
    val compiler = KspKotlinCompiler(tempDir = folder.root)

    val component =
      createSource(
        """
            package test

            import dagger.Module
            import dagger.Provides

            @Module
            class PrimitiveModule {

              @Provides
              fun provideString(): String {
                return "test"
              }
            }
            """
          .trimIndent(),
      )

    val compilation = compiler.compile(component)

    assert(compilation.result.exitCode == KotlinCompilation.ExitCode.COMPILATION_ERROR)
    assert(
      compilation.result.messages.contains("Primitives should be annotated with any qualifier")
    )
  }

  @Test
  fun `primitives can be provided with a qualifier`() {
    val compiler = KspKotlinCompiler(tempDir = folder.root)

    val component =
      createSource(
        """
            package test

            import dagger.Module
            import dagger.Provides
            import javax.inject.Qualifier

            @Qualifier annotation class PrimitiveQualifier

            @Module
            class PrimitiveModule {

              @Provides
              @PrimitiveQualifier
              fun provideString(): String {
                return "test"
              }
            }
            """
          .trimIndent(),
      )

    val compilation = compiler.compile(component)

    assert(compilation.result.exitCode == KotlinCompilation.ExitCode.OK)
  }
}

The second test is our happy path - showing a primitive being provided with a qualifier does not trigger any errors. The first test is demonstrating that an error is thrown with our custom message when a primitive is provided.

The second test we will write is for our BindsValidator plugin:

import com.tschuchort.compiletesting.KotlinCompilation
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
import org.junit.Rule
import org.junit.Test
import org.junit.rules.TemporaryFolder

@OptIn(ExperimentalCompilerApi::class)
class BindsValidatorTest {

  @JvmField @Rule var folder = TemporaryFolder()

  @Test
  fun `items should not be provided if they can be bound instead`() {
    val compiler = KspKotlinCompiler(tempDir = folder.root)

    val component =
      createSource(
        """
        package com.example

        import dagger.Component
        import dagger.Module
        import dagger.Provides
        import javax.inject.Inject
        import javax.inject.Singleton

          interface TestInterface {
            fun hello()
          }

          class TestClass : TestInterface {
            override fun hello() {
              TODO("Not yet implemented")
            }
          }

          @Singleton @Component(modules = [MyModule::class])
          interface MyComponent {
            fun testClass(): TestClass
          }

          @Module object MyModule {
            @Provides
            @Singleton fun providesTestClass() = TestClass()
          }
            """
          .trimIndent(),
      )

    val compilation = compiler.compile(component)

    assert(compilation.result.exitCode == KotlinCompilation.ExitCode.COMPILATION_ERROR)
    assert(
      compilation.result.messages.contains(
        "TestClass should be added to the graph via @Binds instead of @Provides since it extends a single interface"
      )
    )
  }

  @Test
  fun `binds work as intended`() {
    val compiler = KspKotlinCompiler(tempDir = folder.root)

    val component =
      createSource(
        """
        package com.example

        import dagger.Component
        import dagger.Module
        import dagger.Binds
        import javax.inject.Inject
        import javax.inject.Singleton

          interface TestInterface {
            fun hello()
          }

          class TestClass @Inject constructor() : TestInterface {
            override fun hello() {
              TODO("Not yet implemented")
            }
          }

          @Singleton @Component(modules = [MyModule::class])
          interface MyComponent {
            fun testClass(): TestClass
          }

          @Module interface MyModule {
            @Binds
            @Singleton fun bindsTestClass(testClass: TestClass) : TestInterface
          }
            """
          .trimIndent(),
      )

    val compilation = compiler.compile(component)

    assert(compilation.result.exitCode == KotlinCompilation.ExitCode.OK)
  }
}

The second test is our happy path showing that TestClass is successfully bound to TestInterface and no errors are thrown while the first test, which @Provides TestClass instead, shows that our plugin is successfully throwing an error.

Our last test will be for the unique binding plugin.

import com.tschuchort.compiletesting.KotlinCompilation
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
import org.junit.Rule
import org.junit.Test
import org.junit.rules.TemporaryFolder

@OptIn(ExperimentalCompilerApi::class)
class UniqueInjectBindingValidatorTest {

  @JvmField @Rule var folder = TemporaryFolder()

  @Test
  fun `items should not be bound in multiple ways`() {
    val compiler = KspKotlinCompiler(tempDir = folder.root)

    val component =
      createSource(
        """
        package com.example

        import dagger.Component
        import dagger.Module
        import dagger.Provides
        import dagger.Subcomponent
        import javax.inject.Inject
        import javax.inject.Singleton

        @Singleton @Component
        interface MyComponent {
          fun dependency(): Dependency
          fun subcomponent(): MySubcomponent
        }

        @Subcomponent(modules = [MySubmodule::class])
        interface MySubcomponent {
          fun dependency(): Dependency
        }

        @Module object MySubmodule {
          @Provides fun providesDependency() = Dependency("Provides")
        }

        class Dependency @Inject constructor(dep: String)            
          """
          .trimIndent(),
      )

    val compilation = compiler.compile(component)

    assert(compilation.result.exitCode == KotlinCompilation.ExitCode.COMPILATION_ERROR)
    assert(compilation.result.messages.contains("com.example.Dependency is bound multiple times"))
  }

  @Test
  fun `no errors when bound once`() {
    val compiler = KspKotlinCompiler(tempDir = folder.root)

    val component =
      createSource(
        """
        package com.example

        import dagger.Component
        import dagger.Module
        import dagger.Provides
        import javax.inject.Inject
        import javax.inject.Singleton

        @Singleton @Component(modules = [MyModule::class])
        interface MyComponent {
          fun dependency(): Dependency
        }

        @Module object MyModule {
          @Provides fun providesDependency() = Dependency("Provides")
        }

        class Dependency(dep: String)
            """
          .trimIndent(),
      )

    val compilation = compiler.compile(component)

    assert(compilation.result.exitCode == KotlinCompilation.ExitCode.OK)
  }
}

Similarly, the second test is our happy path showing no errors, while the first test shows an error when Dependency is added via both @Inject constructor as well as @Provides .

Consuming Plugins In Your App

Now that we've written our plugins and the accompanying tests, we want to actually consume them in our project.

To do that, we need to add the following to our app's build file:

android {
    defaultConfig {
        ...
        ksp { arg("dagger.fullBindingGraphValidation", "WARNING") }
    }
}

dependencies {
    ksp(project(":dagger-spi-plugins"))
}

After doing so, if we add some test module that violates one of our rules and run ./gradlew :app:kspDebugKotlin, we should see the build fail successfully!

@Module
class PrimitiveModule {
  @Provides
  fun provideString(): String {
    return "test"
  }
}

Please note, the ksp argument dagger.fullBindingGraphValidation and :dagger-spi-plugins dependency need to be added to every module you want the validation to occur in! So if you have a custom plugin setup, you may want to add the dependencies to your library module setup.

Conclusion

Developing custom Dagger SPI plugins has been fun! It not only allowed us to tailor dependency injection to our specific needs but also significantly enhanced our development workflow. By writing tests for our plugins, we ensured their reliability and robustness, providing a safety net that encourages innovation without the fear of regression. Furthermore, our support for Kotlin Symbol Processing (KSP) marked a significant step forward, leveraging the power of Kotlin's tooling to streamline and optimize our build processes.

This is not just about enhancing our tooling; it's about embracing a "shift left" mindset. By moving more checks closer to the compiler and integrating static code analysis tools, we've managed to catch potential issues much earlier in the development cycle. This approach not only saves time and resources but also fosters a culture of quality and responsibility across a team.