1import flask_sqlalchemy as fsa
2
3
4def test_default_query_class(db):
5    class Parent(db.Model):
6        id = db.Column(db.Integer, primary_key=True)
7        children = db.relationship("Child", backref="parent", lazy='dynamic')
8
9    class Child(db.Model):
10        id = db.Column(db.Integer, primary_key=True)
11        parent_id = db.Column(db.Integer, db.ForeignKey('parent.id'))
12
13    p = Parent()
14    c = Child()
15    c.parent = p
16
17    assert type(Parent.query) == fsa.BaseQuery
18    assert type(Child.query) == fsa.BaseQuery
19    assert isinstance(p.children, fsa.BaseQuery)
20    assert isinstance(db.session.query(Parent), fsa.BaseQuery)
21
22
23def test_custom_query_class(app):
24    class CustomQueryClass(fsa.BaseQuery):
25        pass
26
27    db = fsa.SQLAlchemy(app, query_class=CustomQueryClass)
28
29    class Parent(db.Model):
30        id = db.Column(db.Integer, primary_key=True)
31        children = db.relationship("Child", backref="parent", lazy='dynamic')
32
33    class Child(db.Model):
34        id = db.Column(db.Integer, primary_key=True)
35        parent_id = db.Column(db.Integer, db.ForeignKey('parent.id'))
36
37    p = Parent()
38    c = Child()
39    c.parent = p
40
41    assert type(Parent.query) == CustomQueryClass
42    assert type(Child.query) == CustomQueryClass
43    assert isinstance(p.children, CustomQueryClass)
44    assert db.Query == CustomQueryClass
45    assert db.Model.query_class == CustomQueryClass
46    assert isinstance(db.session.query(Parent), CustomQueryClass)
47
48
49def test_dont_override_model_default(app):
50    class CustomQueryClass(fsa.BaseQuery):
51        pass
52
53    db = fsa.SQLAlchemy(app, query_class=CustomQueryClass)
54
55    class SomeModel(db.Model):
56        id = db.Column(db.Integer, primary_key=True)
57        query_class = fsa.BaseQuery
58
59    assert type(SomeModel.query) == fsa.BaseQuery
60