1 files changed, 65 insertions, 0 deletions
diff --git a/kotlinx-coroutines-core/common/test/flow/FlowInvariantsTest.kt b/kotlinx-coroutines-core/common/test/flow/FlowInvariantsTest.kt
index 98406869..e016b031 100644
--- a/kotlinx-coroutines-core/common/test/flow/FlowInvariantsTest.kt
+++ b/kotlinx-coroutines-core/common/test/flow/FlowInvariantsTest.kt
@@ -6,6 +6,7 @@ package kotlinx.coroutines.flow
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.intrinsics.*
 import kotlin.coroutines.*
 import kotlin.reflect.*
 import kotlin.test.*
@@ -214,4 +215,68 @@ class FlowInvariantsTest : TestBase() {
             }
         }
     }
+
+    @Test
+    fun testEmptyCoroutineContext() = runTest {
+        emptyContextTest {
+            map {
+                expect(it)
+                it + 1
+            }
+        }
+    }
+
+    @Test
+    fun testEmptyCoroutineContextTransform() = runTest {
+        emptyContextTest {
+            transform {
+                expect(it)
+                emit(it + 1)
+            }
+        }
+    }
+
+    @Test
+    fun testEmptyCoroutineContextViolation() = runTest {
+        try {
+            emptyContextTest {
+                transform {
+                    expect(it)
+                    kotlinx.coroutines.withContext(Dispatchers.Unconfined) {
+                        emit(it + 1)
+                    }
+                }
+            }
+            expectUnreached()
+        } catch (e: IllegalStateException) {
+            assertTrue(e.message!!.contains("Flow invariant is violated"))
+            finish(2)
+        }
+    }
+
+    private suspend fun emptyContextTest(block: Flow<Int>.() -> Flow<Int>) {
+        suspend fun collector(): Int {
+            var result: Int = -1
+            channelFlow {
+                send(1)
+            }.block()
+                .collect {
+                    expect(it)
+                    result = it
+                }
+            return result
+        }
+
+        val result = runSuspendFun { collector() }
+        assertEquals(2, result)
+        finish(3)
+    }
+
+    private suspend fun runSuspendFun(block: suspend () -> Int): Int {
+        val baseline = Result.failure<Int>(IllegalStateException("Block was suspended"))
+        var result: Result<Int> = baseline
+        block.startCoroutineUnintercepted(Continuation(EmptyCoroutineContext) { result = it })
+        while (result == baseline) yield()
+        return result.getOrThrow()
+    }
 }