package com.fabledsword.minstrel.testutil

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.test.TestDispatcher
import kotlinx.coroutines.test.UnconfinedTestDispatcher
import kotlinx.coroutines.test.resetMain
import kotlinx.coroutines.test.setMain
import org.junit.jupiter.api.extension.AfterEachCallback
import org.junit.jupiter.api.extension.BeforeEachCallback
import org.junit.jupiter.api.extension.ExtensionContext

/**
 * JUnit 5 extension that swaps the kotlinx.coroutines `Dispatchers.Main`
 * for a `TestDispatcher` for the duration of each test, then restores
 * the original. Apply on a test class with:
 *
 *   @ExtendWith(MainDispatcherExtension::class)
 *   class MyViewModelTest { ... }
 *
 * Without this, any ViewModel under test that schedules work on
 * `viewModelScope` (which dispatches to Main by default) will hang or
 * throw "Module with the Main dispatcher had failed to initialize".
 *
 * `UnconfinedTestDispatcher` is the default — coroutines run inline,
 * which makes flow-based tests simpler to assert on.
 */
@OptIn(ExperimentalCoroutinesApi::class)
class MainDispatcherExtension(
    private val testDispatcher: TestDispatcher = UnconfinedTestDispatcher(),
) : BeforeEachCallback, AfterEachCallback {

    override fun beforeEach(context: ExtensionContext) {
        Dispatchers.setMain(testDispatcher)
    }

    override fun afterEach(context: ExtensionContext) {
        Dispatchers.resetMain()
    }
}
