|
| 1 | +from unittest.mock import patch |
| 2 | + |
1 | 3 | import pytest |
2 | 4 | from sqlalchemy import MetaData |
3 | 5 | from sqlalchemy.ext.asyncio import AsyncSession |
@@ -71,3 +73,88 @@ def test_multiple_binds(multiple_config): |
71 | 73 | assert async_bind is not None |
72 | 74 | assert isinstance(sa_manager.get_mapper("async"), registry) |
73 | 75 | assert isinstance(sa_manager.get_session("async"), AsyncSession) |
| 76 | + |
| 77 | + |
| 78 | +async def test_engine_is_disposed_on_cleanup(multiple_config): |
| 79 | + sa_manager = SQLAlchemyBindManager(multiple_config) |
| 80 | + sync_engine = sa_manager.get_bind("default").engine |
| 81 | + async_engine = sa_manager.get_bind("async").engine |
| 82 | + |
| 83 | + original_sync_dispose = sync_engine.dispose |
| 84 | + original_async_dispose = async_engine.dispose |
| 85 | + |
| 86 | + with ( |
| 87 | + patch.object( |
| 88 | + sync_engine, |
| 89 | + "dispose", |
| 90 | + wraps=original_sync_dispose, |
| 91 | + ) as mocked_dispose, |
| 92 | + patch.object( |
| 93 | + type(async_engine), |
| 94 | + "dispose", |
| 95 | + wraps=original_async_dispose, |
| 96 | + ) as mocked_async_dispose, |
| 97 | + ): |
| 98 | + sa_manager = None |
| 99 | + |
| 100 | + mocked_dispose.assert_called_once() |
| 101 | + mocked_async_dispose.assert_called() |
| 102 | + |
| 103 | + |
| 104 | +def test_engine_is_disposed_on_cleanup_even_if_no_loop(multiple_config): |
| 105 | + sa_manager = SQLAlchemyBindManager(multiple_config) |
| 106 | + sync_engine = sa_manager.get_bind("default").engine |
| 107 | + async_engine = sa_manager.get_bind("async").engine |
| 108 | + |
| 109 | + original_sync_dispose = sync_engine.dispose |
| 110 | + original_async_dispose = async_engine.dispose |
| 111 | + |
| 112 | + with ( |
| 113 | + patch.object( |
| 114 | + sync_engine, |
| 115 | + "dispose", |
| 116 | + wraps=original_sync_dispose, |
| 117 | + ) as mocked_dispose, |
| 118 | + patch.object( |
| 119 | + type(async_engine), |
| 120 | + "dispose", |
| 121 | + wraps=original_async_dispose, |
| 122 | + ) as mocked_async_dispose, |
| 123 | + ): |
| 124 | + sa_manager = None |
| 125 | + |
| 126 | + mocked_dispose.assert_called_once() |
| 127 | + mocked_async_dispose.assert_called() |
| 128 | + |
| 129 | + |
| 130 | +def test_engine_is_disposed_on_cleanup_even_if_loop_search_errors_out( |
| 131 | + multiple_config, |
| 132 | +): |
| 133 | + sa_manager = SQLAlchemyBindManager(multiple_config) |
| 134 | + sync_engine = sa_manager.get_bind("default").engine |
| 135 | + async_engine = sa_manager.get_bind("async").engine |
| 136 | + |
| 137 | + original_sync_dispose = sync_engine.dispose |
| 138 | + original_async_dispose = async_engine.dispose |
| 139 | + |
| 140 | + with ( |
| 141 | + patch.object( |
| 142 | + sync_engine, |
| 143 | + "dispose", |
| 144 | + wraps=original_sync_dispose, |
| 145 | + ) as mocked_dispose, |
| 146 | + patch.object( |
| 147 | + type(async_engine), |
| 148 | + "dispose", |
| 149 | + wraps=original_async_dispose, |
| 150 | + ) as mocked_async_dispose, |
| 151 | + patch( |
| 152 | + "asyncio.get_event_loop", |
| 153 | + side_effect=RuntimeError(), |
| 154 | + ) as mocked_get_event_loop, |
| 155 | + ): |
| 156 | + sa_manager = None |
| 157 | + |
| 158 | + mocked_get_event_loop.assert_called_once() |
| 159 | + mocked_dispose.assert_called_once() |
| 160 | + mocked_async_dispose.assert_called() |
0 commit comments