codekingpro/portable-devtools

1# testing/fixtures/orm.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8from __future__ import annotations
9
10from typing import Any
11
12import sqlalchemy as sa
13from .base import TestBase
14from .sql import TablesTest
15from .. import assertions
16from .. import config
17from .. import schema
18from ..entities import BasicEntity
19from ..entities import ComparableEntity
20from ..util import adict
21from ... import orm
22from ...orm import DeclarativeBase
23from ...orm import events as orm_events
24from ...orm import registry
25
26
27class ORMTest(TestBase):
28    @config.fixture
29    def fixture_session(self):
30        return fixture_session()
31
32
33class MappedTest(ORMTest, TablesTest, assertions.AssertsExecutionResults):
34    # 'once', 'each', None
35    run_setup_classes = "once"
36
37    # 'once', 'each', None
38    run_setup_mappers = "each"
39
40    classes: Any = None
41
42    @config.fixture(autouse=True, scope="class")
43    def _setup_tables_test_class(self):
44        cls = self.__class__
45        cls._init_class()
46
47        if cls.classes is None:
48            cls.classes = adict()
49
50        cls._setup_once_tables()
51        cls._setup_once_classes()
52        cls._setup_once_mappers()
53        cls._setup_once_inserts()
54
55        yield
56
57        cls._teardown_once_class()
58        cls._teardown_once_metadata_bind()
59
60    @config.fixture(autouse=True, scope="function")
61    def _setup_tables_test_instance(self):
62        self._setup_each_tables()
63        self._setup_each_classes()
64        self._setup_each_mappers()
65        self._setup_each_inserts()
66
67        yield
68
69        orm.session.close_all_sessions()
70        self._teardown_each_mappers()
71        self._teardown_each_classes()
72        self._teardown_each_tables()
73
74    @classmethod
75    def _teardown_once_class(cls):
76        cls.classes.clear()
77
78    @classmethod
79    def _setup_once_classes(cls):
80        if cls.run_setup_classes == "once":
81            cls._with_register_classes(cls.setup_classes)
82
83    @classmethod
84    def _setup_once_mappers(cls):
85        if cls.run_setup_mappers == "once":
86            cls.mapper_registry, cls.mapper = cls._generate_registry()
87            cls._with_register_classes(cls.setup_mappers)
88
89    def _setup_each_mappers(self):
90        if self.run_setup_mappers != "once":
91            (
92                self.__class__.mapper_registry,
93                self.__class__.mapper,
94            ) = self._generate_registry()
95
96        if self.run_setup_mappers == "each":
97            self._with_register_classes(self.setup_mappers)
98
99    def _setup_each_classes(self):
100        if self.run_setup_classes == "each":
101            self._with_register_classes(self.setup_classes)
102
103    @classmethod
104    def _generate_registry(cls):
105        decl = registry(metadata=cls._tables_metadata)
106        return decl, decl.map_imperatively
107
108    @classmethod
109    def _with_register_classes(cls, fn):
110        """Run a setup method, framing the operation with a Base class
111        that will catch new subclasses to be established within
112        the "classes" registry.
113
114        """
115        cls_registry = cls.classes
116
117        class _Base:
118            def __init_subclass__(cls) -> None:
119                assert cls_registry is not None
120                cls_registry[cls.__name__] = cls
121                super().__init_subclass__()
122
123        class Basic(BasicEntity, _Base):
124            pass
125
126        class Comparable(ComparableEntity, _Base):
127            pass
128
129        cls.Basic = Basic
130        cls.Comparable = Comparable
131        fn()
132
133    def _teardown_each_mappers(self):
134        # some tests create mappers in the test bodies
135        # and will define setup_mappers as None -
136        # clear mappers in any case
137        if self.run_setup_mappers != "once":
138            orm.clear_mappers()
139
140    def _teardown_each_classes(self):
141        if self.run_setup_classes != "once":
142            self.classes.clear()
143
144    @classmethod
145    def setup_classes(cls):
146        pass
147
148    @classmethod
149    def setup_mappers(cls):
150        pass
151
152
153class DeclarativeMappedTest(MappedTest):
154    run_setup_classes = "once"
155    run_setup_mappers = "once"
156
157    @classmethod
158    def _setup_once_tables(cls):
159        pass
160
161    @classmethod
162    def _with_register_classes(cls, fn):
163        cls_registry = cls.classes
164
165        class _DeclBase(DeclarativeBase):
166            __table_cls__ = schema.Table
167            metadata = cls._tables_metadata
168            type_annotation_map = {
169                str: sa.String().with_variant(
170                    sa.String(50), "mysql", "mariadb", "oracle"
171                )
172            }
173
174            def __init_subclass__(cls, **kw) -> None:
175                assert cls_registry is not None
176                cls_registry[cls.__name__] = cls
177                super().__init_subclass__(**kw)
178
179        cls.DeclarativeBasic = _DeclBase
180
181        # sets up cls.Basic which is helpful for things like composite
182        # classes
183        super()._with_register_classes(fn)
184
185        if cls._tables_metadata.tables and cls.run_create_tables:
186            cls._tables_metadata.create_all(config.db)
187
188
189class RemoveORMEventsGlobally:
190    @config.fixture(autouse=True)
191    def _remove_listeners(self):
192        yield
193        orm_events.MapperEvents._clear()
194        orm_events.InstanceEvents._clear()
195        orm_events.SessionEvents._clear()
196        orm_events.InstrumentationEvents._clear()
197        orm_events.QueryEvents._clear()
198
199
200_fixture_sessions = set()
201
202
203def fixture_session(**kw):
204    kw.setdefault("autoflush", True)
205    kw.setdefault("expire_on_commit", True)
206
207    bind = kw.pop("bind", config.db)
208
209    sess = orm.Session(bind, **kw)
210    _fixture_sessions.add(sess)
211    return sess
212
213
214def close_all_sessions():
215    # will close all still-referenced sessions
216    orm.close_all_sessions()
217    _fixture_sessions.clear()
218
219
220def stop_test_class_inside_fixtures(cls):
221    close_all_sessions()
222    orm.clear_mappers()
223
224
225def after_test():
226    if _fixture_sessions:
227        close_all_sessions()
228