Repository: Kotlin/kotlinx.coroutines
File: kotlinx-coroutines-core/jvm/test/ThreadContextElementTest.kt (Kotlin)

Code before:
                assertNull(myThreadLocal.get()) // Asserts value was restored to its origin
            }
        }
    }

    class JobCaptor(val capturees: ArrayList<Job> = ArrayList()) : ThreadContextElement<Unit> {

        companion object Key : CoroutineContext.Key<MyElement>

        override val key: CoroutineContext.Key<*> get() = Key

        override fun updateThreadContext(context: CoroutineContext) {
            capturees.add(context.job)
        }

        override fun restoreThreadContext(context: CoroutineContext, oldState: Unit) {
        }
    }

    /**
     * For stability of the test, it is important to make sure that
     * the parent job actually suspends when calling
     * `withContext(dispatcher2 + CoroutineName("dispatched"))`.
     *
     * Here this requirement is fulfilled by forcing execution on a single thread.
     * However, dispatching is performed with two non-equal dispatchers to simulate multithreaded behavior.
     *
     * Suspend of the parent coroutine [kotlinx.coroutines.DispatchedCoroutine.trySuspend] is out of the control of the test,
     * while being executed concurrently with resume of the child coroutine [kotlinx.coroutines.DispatchedCoroutine.tryResume].
     */
    @Test
    fun testWithContextJobAccess() = runTest {
        val executor = Executors.newSingleThreadExecutor()
        // Emulate non-equal dispatchers
        val executor1 = object : ExecutorService by executor {}
        val executor2 = object : ExecutorService by executor {}
        val dispatcher1 = executor1.asCoroutineDispatcher()
        val dispatcher2 = executor2.asCoroutineDispatcher()
        val captor = JobCaptor()
        val manuallyCaptured = ArrayList<Job>()
        runBlocking(captor + dispatcher1) {
            manuallyCaptured += coroutineContext.job
            withContext(CoroutineName("undispatched")) {
                manuallyCaptured += coroutineContext.job
                // Without forcing of single backing thread the code inside `withContext`
                // may already complete at the moment when the parent coroutine decides
                // whether it needs to suspend or not.
                // If the parent coroutine does not need to suspend, no context capture will be called.
                withContext(dispatcher2 + CoroutineName("dispatched")) {
                    manuallyCaptured += coroutineContext.job
                }

Review comment (suggestion):
You probably mean "to avoid dispatch short-circuit" instead of "simulate multithreaded behavior.", otherwise it's a bit confusing

Code after:
                assertNull(myThreadLocal.get()) // Asserts value was restored to its origin
            }
        }
    }

    class JobCaptor(val capturees: MutableList<String> = CopyOnWriteArrayList()) : ThreadContextElement<Unit> {

        companion object Key : CoroutineContext.Key<MyElement>

        override val key: CoroutineContext.Key<*> get() = Key

        override fun updateThreadContext(context: CoroutineContext) {
            capturees.add("Update: ${context.job}")
        }

        override fun restoreThreadContext(context: CoroutineContext, oldState: Unit) {
            capturees.add("Restore: ${context.job}")
        }
    }

    /**
     * For stability of the test, it is important to make sure that
     * the parent job actually suspends when calling
     * `withContext(dispatcher2 + CoroutineName("dispatched"))`.
     *
     * Here this requirement is fulfilled by forcing execution on a single thread.
     * However, dispatching is performed with two non-equal dispatchers to force dispatching.
     *
     * Suspend of the parent coroutine [kotlinx.coroutines.DispatchedCoroutine.trySuspend] is out of the control of the test,
     * while being executed concurrently with resume of the child coroutine [kotlinx.coroutines.DispatchedCoroutine.tryResume].
     */
    @Test
    fun testWithContextJobAccess() = runTest {
        val executor = Executors.newSingleThreadExecutor()
        // Emulate non-equal dispatchers
        val executor1 = object : ExecutorService by executor {}
        val executor2 = object : ExecutorService by executor {}
        val dispatcher1 = executor1.asCoroutineDispatcher()
        val dispatcher2 = executor2.asCoroutineDispatcher()
        val captor = JobCaptor()
        val manuallyCaptured = mutableListOf<String>()

        fun registerUpdate(job: Job?) = manuallyCaptured.add("Update: $job")
        fun registerRestore(job: Job?) = manuallyCaptured.add("Restore: $job")

        var rootJob: Job? = null
        runBlocking(captor + dispatcher1) {
            rootJob = coroutineContext.job
            registerUpdate(rootJob)
            var undispatchedJob: Job? = null
            withContext(CoroutineName("undispatched")) {