1 
2 /*
3  * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
4  */
5 
6 @file:Suppress("NAMED_ARGUMENTS_NOT_ALLOWED") // KT-22237
7 
8 package kotlinx.coroutines
9 
10 import kotlin.test.*
11 
12 class WithContextTest : TestBase() {
13 
14     @Test
15     fun testThrowException() = runTest {
16         expect(1)
17         try {
18             withContext<Unit>(coroutineContext) {
19                 expect(2)
20                 throw AssertionError()
21             }
22         } catch (e: AssertionError) {
23             expect(3)
24         }
25 
26         yield()
27         finish(4)
28     }
29 
30     @Test
31     fun testThrowExceptionFromWrappedContext() = runTest {
32         expect(1)
33         try {
34             withContext<Unit>(wrapperDispatcher(coroutineContext)) {
35                 expect(2)
36                 throw AssertionError()
37             }
38         } catch (e: AssertionError) {
39             expect(3)
40         }
41 
42         yield()
43         finish(4)
44     }
45 
46     @Test
47     fun testSameContextNoSuspend() = runTest {
48         expect(1)
49         launch(coroutineContext) { // make sure there is not early dispatch here
50             finish(5) // after main exits
51         }
52         expect(2)
53         val result = withContext(coroutineContext) { // same context!
54             expect(3) // still here
55             "OK".wrap()
56         }.unwrap()
57         assertEquals("OK", result)
58         expect(4)
59         // will wait for the first coroutine
60     }
61 
62     @Test
63     fun testSameContextWithSuspend() = runTest {
64         expect(1)
65         launch(coroutineContext) { // make sure there is not early dispatch here
66             expect(4)
67         }
68         expect(2)
69         val result = withContext(coroutineContext) { // same context!
70             expect(3) // still here
71             yield() // now yields to launch!
72             expect(5)
73             "OK".wrap()
74         }.unwrap()
75         assertEquals("OK", result)
76         finish(6)
77     }
78 
79     @Test
80     fun testCancelWithJobNoSuspend() = runTest {
81         expect(1)
82         launch(coroutineContext) { // make sure there is not early dispatch to here
83             finish(6) // after main exits
84         }
85         expect(2)
86         val job = Job()
87         try {
88             withContext(coroutineContext + job) {
89                 // same context + new job
90                 expect(3) // still here
91                 job.cancel() // cancel out job!
92                 try {
93                     yield() // shall throw CancellationException
94                     expectUnreached()
95                 } catch (e: CancellationException) {
96                     expect(4)
97                 }
98                 "OK".wrap()
99             }
100 
101             expectUnreached()
102         } catch (e: CancellationException) {
103             expect(5)
104             // will wait for the first coroutine
105         }
106     }
107 
108     @Test
109     fun testCancelWithJobWithSuspend() = runTest(
110         expected = { it is CancellationException }
111     ) {
112         expect(1)
113         launch(coroutineContext) { // make sure there is not early dispatch to here
114             expect(4)
115         }
116         expect(2)
117         val job = Job()
118         withContext(coroutineContext + job) { // same context + new job
119             expect(3) // still here
120             yield() // now yields to launch!
121             expect(5)
122             job.cancel() // cancel out job!
123             try {
124                 yield() // shall throw CancellationException
125                 expectUnreached()
126             } catch (e: CancellationException) {
127                 finish(6)
128             }
129             "OK".wrap()
130         }
131         // still fails, because parent job was cancelled
132         expectUnreached()
133     }
134 
135     @Test
136     fun testRunCancellableDefault() = runTest(
137         expected = { it is CancellationException }
138     ) {
139         val job = Job()
140         job.cancel() // cancel before it has a chance to run
141         withContext(job + wrapperDispatcher(coroutineContext)) {
142             expectUnreached() // will get cancelled
143         }
144     }
145 
146     @Test
147     fun testRunCancellationUndispatchedVsException() = runTest {
148         expect(1)
149         var job: Job? = null
150         job = launch(start = CoroutineStart.UNDISPATCHED) {
151             expect(2)
152             try {
153                 // Same dispatcher, different context
154                 withContext<Unit>(CoroutineName("testRunCancellationUndispatchedVsException")) {
155                     expect(3)
156                     yield() // must suspend
157                     expect(5)
158                     job!!.cancel() // cancel this job _before_ it throws
159                     throw TestException()
160                 }
161             } catch (e: TestException) {
162                 // must have caught TextException
163                 expect(6)
164             }
165         }
166         expect(4)
167         yield() // to coroutineScope
168         finish(7)
169     }
170 
171     @Test
172     fun testRunCancellationDispatchedVsException() = runTest {
173         expect(1)
174         var job: Job? = null
175         job = launch(start = CoroutineStart.UNDISPATCHED) {
176             expect(2)
177             try {
178                 // "Different" dispatcher (schedules execution back and forth)
179                 withContext<Unit>(wrapperDispatcher(coroutineContext)) {
180                     expect(4)
181                     yield() // must suspend
182                     expect(6)
183                     job!!.cancel() // cancel this job _before_ it throws
184                     throw TestException()
185                 }
186             } catch (e: TestException) {
187                 // must have caught TextException
188                 expect(8)
189             }
190         }
191         expect(3)
192         yield() // withContext is next
193         expect(5)
194         yield() // withContext again
195         expect(7)
196         yield() // to catch block
197         finish(9)
198     }
199 
200     @Test
201     fun testRunSelfCancellationWithException() = runTest {
202         expect(1)
203         var job: Job? = null
204         job = launch(Job()) {
205             try {
206                 expect(3)
207                 withContext<Unit>(wrapperDispatcher(coroutineContext)) {
208                     require(isActive)
209                     expect(5)
210                     job!!.cancel()
211                     require(!isActive)
212                     throw TestException() // but throw an exception
213                 }
214             } catch (e: Throwable) {
215                 expect(7)
216                 // make sure TestException, not CancellationException is thrown
217                 assertTrue(e is TestException, "Caught $e")
218             }
219         }
220         expect(2)
221         yield() // to the launched job
222         expect(4)
223         yield() // again to the job
224         expect(6)
225         yield() // again to exception handler
226         finish(8)
227     }
228 
229     @Test
230     fun testRunSelfCancellation() = runTest {
231         expect(1)
232         var job: Job? = null
233         job = launch(Job()) {
234             try {
235                 expect(3)
236                 withContext(wrapperDispatcher(coroutineContext)) {
237                     require(isActive)
238                     expect(5)
239                     job!!.cancel() // cancel itself
240                     require(!isActive)
241                     "OK".wrap()
242                 }
243                 expectUnreached()
244             } catch (e: Throwable) {
245                 expect(7)
246                 // make sure CancellationException is thrown
247                 assertTrue(e is CancellationException, "Caught $e")
248             }
249         }
250 
251         expect(2)
252         yield() // to the launched job
253         expect(4)
254         yield() // again to the job
255         expect(6)
256         yield() // again to exception handler
257         finish(8)
258     }
259 
260     @Test
261     fun testWithContextScopeFailure() = runTest {
262         expect(1)
263         try {
264             withContext(wrapperDispatcher(coroutineContext)) {
265                 expect(2)
266                 // launch a child that fails
267                 launch {
268                     expect(4)
269                     throw TestException()
270                 }
271                 expect(3)
272                 "OK".wrap()
273             }
274             expectUnreached()
275         } catch (e: TestException) {
276             // ensure that we can catch exception outside of the scope
277             expect(5)
278         }
279         finish(6)
280     }
281 
282     @Test
283     fun testWithContextChildWaitSameContext() = runTest {
284         expect(1)
285         withContext(coroutineContext) {
286             expect(2)
287             launch {
288                 // ^^^ schedules to main thread
289                 expect(4) // waits before return
290             }
291             expect(3)
292             "OK".wrap()
293         }.unwrap()
294         finish(5)
295     }
296 
297     @Test
298     fun testWithContextChildWaitWrappedContext() = runTest {
299         expect(1)
300         withContext(wrapperDispatcher(coroutineContext)) {
301             expect(2)
302             launch {
303                 // ^^^ schedules to main thread
304                 expect(4) // waits before return
305             }
306             expect(3)
307             "OK".wrap()
308         }.unwrap()
309         finish(5)
310     }
311 
312     @Test
313     fun testIncompleteWithContextState() = runTest {
314         lateinit var ctxJob: Job
315         withContext(wrapperDispatcher(coroutineContext)) {
316             ctxJob = coroutineContext[Job]!!
317             ctxJob.invokeOnCompletion { }
318         }
319 
320         ctxJob.join()
321         assertTrue(ctxJob.isCompleted)
322         assertFalse(ctxJob.isActive)
323         assertFalse(ctxJob.isCancelled)
324     }
325 
326     @Test
327     fun testWithContextCancelledJob() = runTest {
328         expect(1)
329         val job = Job()
330         job.cancel()
331         try {
332             withContext(job) {
333                 expectUnreached()
334             }
335         } catch (e: CancellationException) {
336             expect(2)
337         }
338         finish(3)
339     }
340 
341     @Test
342     fun testWithContextCancelledThisJob() = runTest(
343         expected = { it is CancellationException }
344     ) {
345         coroutineContext.cancel()
346         withContext(wrapperDispatcher(coroutineContext)) {
347             expectUnreached()
348         }
349         expectUnreached()
350     }
351 
352     @Test
353     fun testSequentialCancellation() = runTest {
354         val job = launch {
355             expect(1)
356             withContext(wrapperDispatcher()) {
357                 expect(2)
358             }
359             expectUnreached()
360         }
361 
362         yield()
363         val job2 = launch {
364             expect(3)
365             job.cancel()
366         }
367 
368         joinAll(job, job2)
369         finish(4)
370     }
371 
372     private class Wrapper(val value: String) : Incomplete {
373         override val isActive: Boolean
374             get() =  error("")
375         override val list: NodeList?
376             get() = error("")
377     }
378 
379     private fun String.wrap() = Wrapper(this)
380     private fun Wrapper.unwrap() = value
381 }
382