【Kotlin】CancellationExceptionをテストする

やったこと

　例外処理とCancellationExceptiontonの扱いを保証するUnitTestを書いたので備忘のためメモ

CancellationExceptionの扱いが面倒

　DataA、DataBというデータから作られるクラスDataCがあるとします↓

data class DataC(...) {
    constructor(
        dataA: DataA,
        dataB: DataB,
    ) : this(...)
}

　DataA、DataBはそれぞれDataSourceA.loadA()、DataSourceB.loadB()というsuspend関数から取得できるとします。
　失敗をどう扱うかは色々だと思いますが、今回は簡単のためにそれぞれの関数は結果を返すか、例外を吐くかのどちらかの動作をするとします↓

// loadA()は失敗したら例外を吐く
interface DataSourceA {
    suspend fun loadA(): DataA
}

// loadB()は失敗したら例外を吐く
interface DataSourceB {
    suspend fun loadB(): DataB
}

　loadA()、loadB()が独立した処理のとき、asyncで両方の関数を非同期で呼びDataCを作成する && どちらかが失敗したら、もう片方を待たずに失敗を返す関数RepositoryC.loadC()を考えます。
　ここも返り値は様々あるかと思いますが、今回はkotlin.Resultを使っています↓

interface RepositoryC {
    suspend fun loadC(): Result<DataC>
}

class DefaultRepositoryC(
    private val dataSourceA: DataSourceA,
    private val dataSourceB: DataSourceB,
    private val dispatcher: CoroutineDispatcher = Dispatchers.IO,
) : RepositoryC {
    override suspend fun loadC(): Result<DataC> = runCatching {
        withContext(dispatcher) {
            val deferredA = async { dataSourceA.loadA() }
            val deferredB = async { dataSourceB.loadB() }
            DataC(deferredA.await(), deferredB.await())
        }
    }.onFailure { if (it is CancellationException) throw it }
}

　.onFailure {}で失敗時の例外がCancellationExceptionの場合だけ例外を投げ返しています。
　これはKotlinのコルーチンがCancellationExceptionを利用してジョブのキャンセルを伝えるため、この例外をkotlin.Resultに詰めて返してしまうと、キャンセル時に正常にキャンセルが伝播しなくなるためです↓

Caution: To enable coroutine cancellation, don't consume exceptions of type CancellationException (don't catch them, or always rethrow them if caught). Prefer catching specific exception types like IOException over generic types like Exception or Throwable.

　当然気を付けていれば問題はないのですが、正直かなりややこしいです。
　例えばこのloadC()のrunCatching {}とwithContext(dispatcher) {}の順番を入れ替えると、キャンセル以外の失敗時にも例外が親コルーチンまで返ります↓

NG

    override suspend fun loadC(): Result<DataC> = withContext(dispatcher) { // async内部で例外を起こすとcatchを超えてここに返る
        runCatching {
            val deferredA = async { DataSourceA.loadA() } // loadAは例外を吐く
            val deferredB = async { DataSourceB.loadB() } // loadBは例外を吐く
            DataC(deferredA.await(), deferredB.await())
        }.onFailure { if (it is CancellationException) throw it }
    }

　これではおちおち触れないため、このコルーチンと例外周りのロジックを保証する用のUnitTestを用意しました。

UnitTestで保証する

　JUnit4、org.jetbrains.kotlinx:kotlinx-coroutines-test、io.mockk:mockkを使用しています。

import io.mockk.coEvery
import io.mockk.junit4.MockKRule
import io.mockk.mockk
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.StandardTestDispatcher
import kotlinx.coroutines.test.TestScope
import kotlinx.coroutines.test.advanceTimeBy
import kotlinx.coroutines.test.runTest
import org.junit.Assert.assertTrue
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import kotlin.coroutines.cancellation.CancellationException

@OptIn(ExperimentalCoroutinesApi::class)
class RepositoryCTest {
    @get:Rule
    val mockkRule = MockKRule(this)
    
    private val dataSourceA: DataSourceA = mockk()
    private val dataSourceB: DataSourceB = mockk()
    
    @Before
    fun setup() {
        coEvery { dataSourceB.loadB() } answers { mockk(relaxed = true) }
    }

    // 正常系
    @Test
    fun `loadC success`() = runTest {
        val repository = createRepository()
        coEvery { dataSourceA.loadA() } answers { mockk(relaxed = true) }
        val result = repository.loadC()
        assertTrue(result.isSuccess)
    }

    // loadAの実行に失敗
    @Test
    fun `loadC error`() = runTest {
        val repository = createRepository()
        coEvery { dataSourceA.loadA() } throws (Exception(""))
        val result = repository.loadC()
        assertTrue(result.isFailure)
    }

    // 親コルーチンがキャンセル
    @Test
    fun `loadC parent cancel`() = runTest {
        val repository = createRepository()
        coEvery { dataSourceA.loadA() } coAnswers {
            delay(1000)
            mockk(relaxed = true)
        }
        
        var isJobContinue = false
        val job = launch {
            repository.loadC()
            isJobContinue = true
        }
        advanceTimeBy(500)
        job.cancelAndJoin()
        assertTrue(!isJobContinue)
    }

    // 子コルーチンからキャンセル通知
    @Test
    fun `loadC child cancel`() = runTest {
        val repository = createRepository()
        coEvery { dataSourceA.loadA() } throws (CancellationException(""))
        
        var isJobContinue = false
        val job = launch {
            repository.loadC()
            isJobContinue = true
        }
        job.join()
        assertTrue(!isJobContinue)
    }
    
    private fun TestScope.createRepository() = DefaultRepositoryC(
        dataSourceA = dataSourceA,
        dataSourceB = dataSourceB,
        dispatcher = StandardTestDispatcher(testScheduler),
    )
}

感想

やはりテストは大事だなと思いました。