Use different types for TestState

mirai-core/src/commonMain/kotlin/network/handler/NetworkHandlerSupport.kt

@@ -162,6 +162,12 @@ internal abstract class NetworkHandlerSupport(
      *
      * State can only be changed inside [setState].
      *
+     * **IMPORTANT implementation notes:**
+     *
+     * You must create subclasses of [BaseStateImpl] for EVERY SINGLE [NetworkHandler.State].
+     * **DO NOT** use same type for more than one [NetworkHandler.State],
+     * otherwise [setState] will refuse updating state in some concurrent situations and will be very difficult to debug.
+     *
      * **IMPORTANT notes to lifecycle:**
      *
      * Normally if the state is set to [NetworkHandler.State.CLOSED] by [setState], [selector][NetworkHandlerSelector] may reinitialize an instance.
@@ -177,8 +183,11 @@ internal abstract class NetworkHandlerSupport(
         final override val coroutineContext: CoroutineContext =
             [email protected] + Job([email protected])
 
+        // Important: read the above doc before implementing BaseStateImpl.
+
         // Do not use init blocks to launch anything. Do use [startState]
 
+
         /**
          * Starts things that should be done in this state.
          *
@@ -271,16 +280,14 @@ internal abstract class NetworkHandlerSupport(
     internal val lockForSetStateWithOldInstance = SynchronizedObject()
 
     /**
-     * This can only be called by [setState] or in tests.
-     *
-     * [newType] can be `null` **iff in tests**, to ignore checks.
+     * This can only be called by [setState] or in tests. Note:
      */
     //
     @TestOnly
-    internal fun <S : BaseStateImpl> setStateImpl(newType: KClass<S>?, new: () -> S): S? =
+    internal fun <S : BaseStateImpl> setStateImpl(newType: KClass<S>, new: () -> S): S? =
         lock.withLock {
             val old = _state
-            if (newType != null && old::class == newType) return@withLock null // already set to expected state by another thread. Avoid replications.
+            if (old::class == newType) return@withLock null // already set to expected state by another thread. Avoid replications.
             if (old.correspondingState == NetworkHandler.State.CLOSED) return@withLock null // CLOSED is final.
 
             val stateObserver = context.getOrNull(StateObserver)

mirai-core/src/commonTest/kotlin/network/framework/TestNetworkHandler.kt

@@ -15,6 +15,7 @@ import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import net.mamoe.mirai.internal.QQAndroidBot
 import net.mamoe.mirai.internal.network.handler.NetworkHandler
+import net.mamoe.mirai.internal.network.handler.NetworkHandler.State.*
 import net.mamoe.mirai.internal.network.handler.NetworkHandlerContext
 import net.mamoe.mirai.internal.network.handler.NetworkHandlerSupport
 import net.mamoe.mirai.internal.network.handler.logger
@@ -32,7 +33,7 @@ internal open class TestNetworkHandler(
     class Connection
 
     @Suppress("EXPOSED_SUPER_CLASS")
-    internal open inner class TestState(
+    internal abstract inner class TestState(
         correspondingState: NetworkHandler.State
     ) : BaseStateImpl(correspondingState) {
         val resumeDeferred = CompletableDeferred<Unit>()
@@ -44,8 +45,8 @@ internal open class TestNetworkHandler(
             resumeCount.incrementAndGet()
             resumeDeferred.complete(Unit)
             when (this.correspondingState) {
-                NetworkHandler.State.INITIALIZED -> {
-                    setState(NetworkHandler.State.CONNECTING)
+                INITIALIZED -> {
+                    setState(CONNECTING)
                 }
                 else -> {
                 }
@@ -57,13 +58,26 @@ internal open class TestNetworkHandler(
         }
     }
 
+    internal inner class TestStateInitial : TestState(INITIALIZED)
+    internal inner class TestStateConnecting : TestState(CONNECTING)
+    internal inner class TestStateLoading : TestState(LOADING)
+    internal inner class TestStateOK : TestState(OK)
+    internal inner class TestStateClosed : TestState(CLOSED)
+
     @OptIn(TestOnly::class)
     fun setState(correspondingState: NetworkHandler.State): TestState? {
         // `null` means ignore checks. All test states have same type TestState.
-        return setStateImpl(null) { TestState(correspondingState) }
+        val state: TestState = when (correspondingState) {
+            INITIALIZED -> TestStateInitial()
+            CONNECTING -> TestStateConnecting()
+            LOADING -> TestStateLoading()
+            OK -> TestStateOK()
+            CLOSED -> TestStateClosed()
+        }
+        return setStateImpl(TestState::class) { state }
     }
 
-    private val initialState = TestState(NetworkHandler.State.INITIALIZED)
+    private val initialState = TestStateInitial()
     override fun initialState(): BaseStateImpl = initialState
 
     val sendPacket get() = ConcurrentLinkedQueue<OutgoingPacket>()
@@ -75,19 +89,19 @@ internal open class TestNetworkHandler(
 
 
     override fun setStateClosed(exception: Throwable?): TestState? {
-        return setState(NetworkHandler.State.CLOSED)
+        return setState(CLOSED)
     }
 
     override fun setStateConnecting(exception: Throwable?): TestState? {
-        return setState(NetworkHandler.State.CONNECTING)
+        return setState(CONNECTING)
     }
 
     override fun setStateOK(conn: Connection, exception: Throwable?): TestState? {
         exception?.printStackTrace()
-        return setState(NetworkHandler.State.OK)
+        return setState(OK)
     }
 
     override fun setStateLoading(conn: Connection): TestState? {
-        return setState(NetworkHandler.State.LOADING)
+        return setState(LOADING)
     }
 }