Add tests for query_entities

tests/functions/test_query_entities.py

@@ -0,0 +1,70 @@
+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]