forked from pallets-eco/flask-sqlalchemy
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_extension_object.py
More file actions
75 lines (56 loc) · 2.56 KB
/
test_extension_object.py
File metadata and controls
75 lines (56 loc) · 2.56 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
from __future__ import annotations
import typing as t
import pytest
import sqlalchemy as sa
import sqlalchemy.orm as sa_orm
from flask import Flask
from werkzeug.exceptions import NotFound
from flask_sqlalchemy import SQLAlchemy
from flask_sqlalchemy.record_queries import get_recorded_queries
@pytest.mark.usefixtures("app_ctx")
def test_get_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
item = Todo()
db.session.add(item)
db.session.commit()
assert db.get_or_404(Todo, 1) is item
with pytest.raises(NotFound):
db.get_or_404(Todo, 2)
def test_get_or_404_kwargs(app: Flask, model_class: t.Any) -> None:
app.config["SQLALCHEMY_RECORD_QUERIES"] = True
db = SQLAlchemy(app, model_class=model_class)
if issubclass(db.Model, (sa_orm.MappedAsDataclass)):
class User(db.Model):
id: sa_orm.Mapped[int] = sa_orm.mapped_column(
sa.Integer, primary_key=True, init=False
)
class Todo(db.Model):
id: sa_orm.Mapped[int] = sa_orm.mapped_column(
sa.Integer, primary_key=True, init=False
)
user_id: sa_orm.Mapped[int] = sa_orm.mapped_column(
sa.ForeignKey(User.id), init=False
)
user: sa_orm.Mapped[User] = sa_orm.relationship(User)
elif issubclass(db.Model, (sa_orm.DeclarativeBase, sa_orm.DeclarativeBaseNoMeta)):
class User(db.Model): # type: ignore[no-redef]
id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True)
class Todo(db.Model): # type: ignore[no-redef]
id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True)
user_id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.ForeignKey(User.id))
user: sa_orm.Mapped[User] = sa_orm.relationship(User)
else:
class User(db.Model): # type: ignore[no-redef]
id = sa.Column(db.Integer, primary_key=True) # type: ignore[var-annotated]
class Todo(db.Model): # type: ignore[no-redef]
id = sa.Column(sa.Integer, primary_key=True)
user_id = sa.Column(sa.ForeignKey(User.id)) # type: ignore[var-annotated]
user = db.relationship(User)
with app.app_context():
db.create_all()
db.session.add(Todo(user=User()))
db.session.commit()
with app.app_context():
item = db.get_or_404(Todo, 1, options=[db.joinedload(Todo.user)])
assert item.user.id == 1
# one query with join, no second query when accessing relationship
assert len(get_recorded_queries()) == 1