75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
import sqlalchemy as sa
|
|
from tests import TestCase
|
|
from sqlalchemy_utils.functions import query_entities
|
|
|
|
|
|
class TestQueryEntities(TestCase):
|
|
def create_models(self):
|
|
class User(self.Base):
|
|
__tablename__ = 'user'
|
|
id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
|
|
name = sa.Column(sa.Unicode(255))
|
|
|
|
class Category(self.Base):
|
|
__tablename__ = 'category'
|
|
id = sa.Column(sa.Integer, primary_key=True)
|
|
name = sa.Column(sa.Unicode(255))
|
|
|
|
class Article(self.Base):
|
|
__tablename__ = 'article'
|
|
id = sa.Column(sa.Integer, primary_key=True)
|
|
name = sa.Column(sa.Unicode(255))
|
|
author_id = sa.Column(
|
|
sa.Integer, sa.ForeignKey(User.id), index=True
|
|
)
|
|
category_id = sa.Column(sa.Integer, sa.ForeignKey(Category.id))
|
|
|
|
category = sa.orm.relationship(
|
|
Category,
|
|
primaryjoin=category_id == Category.id,
|
|
backref=sa.orm.backref(
|
|
'articles',
|
|
)
|
|
)
|
|
|
|
self.User = User
|
|
self.Category = Category
|
|
self.Article = Article
|
|
|
|
def test_simple_query(self):
|
|
query = self.session.query(self.User)
|
|
assert list(query_entities(query)) == [self.User]
|
|
|
|
def test_column_entity(self):
|
|
query = self.session.query(self.User.id)
|
|
assert list(query_entities(query)) == [self.User]
|
|
|
|
def test_column_entity_with_label(self):
|
|
query = self.session.query(self.User.id.label('id'))
|
|
assert list(query_entities(query)) == [self.User]
|
|
|
|
def test_with_subquery(self):
|
|
number_of_sales = (
|
|
sa.select(
|
|
[sa.func.count(self.Article.id)],
|
|
)
|
|
.select_from(
|
|
self.Article.__table__
|
|
)
|
|
).label('number_of_articles')
|
|
|
|
query = self.session.query(self.User, number_of_sales)
|
|
assert list(query_entities(query)) == [self.User]
|
|
|
|
def test_mapper(self):
|
|
query = self.session.query(self.User.__mapper__)
|
|
assert list(query_entities(query)) == [self.User]
|
|
|
|
def test_joins(self):
|
|
query = self.session.query(self.User.__mapper__).join(self.Article)
|
|
assert list(query_entities(query)) == [self.User, self.Article]
|
|
|
|
def test_aliased_entity(self):
|
|
query = self.session.query(sa.orm.aliased(self.User))
|
|
assert list(query_entities(query)) == [self.User]
|