|
1 | 1 | import asyncio |
| 2 | +import contextvars |
2 | 3 | import functools |
3 | 4 | import multiprocessing |
4 | 5 | import sys |
|
13 | 14 | import pytest |
14 | 15 |
|
15 | 16 | from asgiref.sync import ( |
| 17 | + AsyncSingleThreadContext, |
16 | 18 | ThreadSensitiveContext, |
17 | 19 | async_to_sync, |
18 | 20 | iscoroutinefunction, |
@@ -544,6 +546,98 @@ def inner(result): |
544 | 546 | assert result_1["thread"] == result_2["thread"] |
545 | 547 |
|
546 | 548 |
|
| 549 | +def test_async_single_thread_context_matches(): |
| 550 | + """ |
| 551 | + Tests that functions wrapped with async_to_sync and executed within an |
| 552 | + AsyncSingleThreadContext run on the same thread, even without a main_event_loop. |
| 553 | + """ |
| 554 | + result_1 = {} |
| 555 | + result_2 = {} |
| 556 | + |
| 557 | + async def store_thread_async(result): |
| 558 | + result["thread"] = threading.current_thread() |
| 559 | + |
| 560 | + with AsyncSingleThreadContext(): |
| 561 | + async_to_sync(store_thread_async)(result_1) |
| 562 | + async_to_sync(store_thread_async)(result_2) |
| 563 | + |
| 564 | + # They should not have run in the main thread, and on the same threads |
| 565 | + assert result_1["thread"] != threading.current_thread() |
| 566 | + assert result_1["thread"] == result_2["thread"] |
| 567 | + |
| 568 | + |
| 569 | +def test_async_single_thread_nested_context(): |
| 570 | + """ |
| 571 | + Tests that behavior remains the same when using nested context managers. |
| 572 | + """ |
| 573 | + result_1 = {} |
| 574 | + result_2 = {} |
| 575 | + |
| 576 | + @async_to_sync |
| 577 | + async def store_thread(result): |
| 578 | + result["thread"] = threading.current_thread() |
| 579 | + |
| 580 | + with AsyncSingleThreadContext(): |
| 581 | + store_thread(result_1) |
| 582 | + |
| 583 | + with AsyncSingleThreadContext(): |
| 584 | + store_thread(result_2) |
| 585 | + |
| 586 | + # They should not have run in the main thread, and on the same threads |
| 587 | + assert result_1["thread"] != threading.current_thread() |
| 588 | + assert result_1["thread"] == result_2["thread"] |
| 589 | + |
| 590 | + |
| 591 | +def test_async_single_thread_context_without_async_work(): |
| 592 | + """ |
| 593 | + Tests everything works correctly without any async_to_sync calls. |
| 594 | + """ |
| 595 | + with AsyncSingleThreadContext(): |
| 596 | + pass |
| 597 | + |
| 598 | + |
| 599 | +def test_async_single_thread_context_success_share_context(): |
| 600 | + """ |
| 601 | + Tests that we share context between different async_to_sync functions. |
| 602 | + """ |
| 603 | + connection = contextvars.ContextVar("connection") |
| 604 | + connection.set(0) |
| 605 | + |
| 606 | + async def handler(): |
| 607 | + connection.set(connection.get(0) + 1) |
| 608 | + |
| 609 | + with AsyncSingleThreadContext(): |
| 610 | + async_to_sync(handler)() |
| 611 | + async_to_sync(handler)() |
| 612 | + |
| 613 | + assert connection.get() == 2 |
| 614 | + |
| 615 | + |
| 616 | +@pytest.mark.asyncio |
| 617 | +async def test_async_single_thread_context_matches_from_async_thread(): |
| 618 | + """ |
| 619 | + Tests that we use main_event_loop for running async_to_sync functions executed |
| 620 | + within an AsyncSingleThreadContext. |
| 621 | + """ |
| 622 | + result_1 = {} |
| 623 | + result_2 = {} |
| 624 | + |
| 625 | + @async_to_sync |
| 626 | + async def store_thread_async(result): |
| 627 | + result["thread"] = threading.current_thread() |
| 628 | + |
| 629 | + def inner(): |
| 630 | + with AsyncSingleThreadContext(): |
| 631 | + store_thread_async(result_1) |
| 632 | + store_thread_async(result_2) |
| 633 | + |
| 634 | + await sync_to_async(inner)() |
| 635 | + |
| 636 | + # They should both have run in the current thread. |
| 637 | + assert result_1["thread"] == threading.current_thread() |
| 638 | + assert result_1["thread"] == result_2["thread"] |
| 639 | + |
| 640 | + |
547 | 641 | @pytest.mark.asyncio |
548 | 642 | async def test_thread_sensitive_with_context_matches(): |
549 | 643 | result_1 = {} |
|
0 commit comments