package kotlinx.coroutines.flow

import kotlinx.coroutines.*
import kotlin.test.*

class FlattenMergeTest : FlatMapMergeBaseTest() {

    override fun <T> Flow<T>.flatMap(mapper: suspend (T) -> Flow<T>): Flow<T> = map(mapper).flattenMerge()

    @Test
    override fun testFlatMapConcurrency() = runTest {
        var concurrentRequests = 0
        val flow = (1..100).asFlow().map { value ->
            flow {
                ++concurrentRequests
                emit(value)
                delay(Long.MAX_VALUE)
            }
        }.flattenMerge(concurrency = 2)

        val consumer = launch {
            flow.collect { value ->
                expect(value)
            }
        }

        repeat(4) {
            yield()
        }

        assertEquals(2, concurrentRequests)
        consumer.cancelAndJoin()
        finish(3)
    }

    @Test
    fun testContextPreservationAcrossFlows() = runTest {
        val result = flow {
            flowOf(1, 2).flatMapMerge {
                flow {
                    yield()
                    emit(it)
                }
            }.collect {
                emit(it)
            }
        }.toList()
        assertEquals(listOf(1, 2), result)
    }
}