WorkManagerTestRule implementation

WorkManagerTestRule.kt

import android.content.Context
import androidx.test.platform.app.InstrumentationRegistry
import androidx.work.Configuration
import androidx.work.Operation
import androidx.work.WorkManager
import androidx.work.WorkQuery
import androidx.work.WorkRequest
import androidx.work.testing.SynchronousExecutor
import androidx.work.testing.TestDriver
import androidx.work.testing.WorkManagerTestInitHelper
import java.util.UUID
import org.junit.rules.ExternalResource

/**
 * A JUnit rule that initializes WorkManager for tests. 
 *
 * For use with android instrumentation or robolectric tests involving WorkManager
 */
class WorkManagerTestRule(
  private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
) : ExternalResource() {

  private val executor = SynchronousExecutor()

  /**
   * [Configuration.Builder] for the [WorkManager] used in the test.
   *
   * The default value sets the executor to be the [SynchronousExecutor]. To add more state or
   * change the builder, call this before your test setup returns.
   */
  var configBuilder: Configuration.Builder =
    Configuration.Builder().setExecutor(executor).setTaskExecutor(executor)

  val driver: TestDriver? by lazy {
    WorkManagerTestInitHelper.getTestDriver(context)
  }

  val workManager by lazy {
    WorkManager.getInstance(context)
  }

  override fun before() {
    super.before()
    WorkManagerTestInitHelper.initializeTestWorkManager(context, configBuilder.build())
  }

  override fun after() {
    super.after()
    WorkManagerTestInitHelper.closeWorkDatabase()
  }
}

/**
 * Enqueues a [WorkRequest]. The work request only runs if there are no constraints or all the
 * constraints are met.
 *
 * Shortcut for [WorkManager.enqueue]
 */
fun WorkManagerTestRule.enqueue(request: WorkRequest): Operation {
  return workManager.enqueue(request)
}

/**
 * Enqueues a list of [WorkRequest]s. The work requests only run if there are no constraints or all
 * the constraints are met.
 *
 * Shortcut for [WorkManager.enqueue]
 */
fun WorkManagerTestRule.enqueue(requests: List<WorkRequest>): Operation {
  return workManager.enqueue(requests)
}

/**
 * Enqueues a [WorkRequest], and then meets its constraints to execute it.
 *
 * Shortcut for [WorkManager.enqueue] and [TestDriver.setAllConstraintsMet]
 */
fun WorkManagerTestRule.execute(request: WorkRequest) {
  with(request) {
    workManager.enqueue(this)
    driver?.setAllConstraintsMet(id)
  }
}

/**
 * Enqueues a list of [WorkRequest]s, and then meets all their constraints to execute them.
 *
 * Shortcut for [WorkManager.enqueue] and [TestDriver.setAllConstraintsMet]
 */
fun WorkManagerTestRule.execute(requests: List<WorkRequest>) {
  requests.forEach { request ->
    execute(request)
  }
}

/** Executes all work that match the given [query] by meeting all their constraints to. */
fun WorkManagerTestRule.execute(query: WorkQuery) {
  val infos = workManager.getWorkInfos(query).get()
  infos.forEach { setAllConstraintsMet(it.id) }
}

/**
 * Sets all constraints on the WorkManager work with the given [workSpecId]. Shortcut for
 * [TestDriver.setAllConstraintsMet]
 */
fun WorkManagerTestRule.setAllConstraintsMet(workSpecId: UUID) {
  driver?.setAllConstraintsMet(workSpecId)
}

/** Sets all constraints on all the WorkManager works that match the given [query]. */
fun WorkManagerTestRule.setAllConstraintsMet(query: WorkQuery) {
  val infos = workManager.getWorkInfos(query).get()
  infos.forEach { setAllConstraintsMet(it.id) }
}

/** Returns the [TestDriver] for the [WorkManagerTestRule]. */
fun WorkManagerTestRule.requireDriver() =
  requireNotNull(driver) {
    "Driver is not initialized. Ensure that you have not called this function before your @Before" +
      " setup has been called."
  }