Test annotate fix.

coleifer · Feb 10, 2015 · 1091113 · 1091113
1 parent 5aea8f6
commit 1091113
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/peewee.py b/peewee.py
@@ -2465,7 +2465,10 @@ def alias(self, alias=None):
     def annotate(self, rel_model, annotation=None):
         if annotation is None:
             annotation = fn.Count(rel_model._meta.primary_key).alias('count')
-        query = self.clone()
+        if self._query_ctx == rel_model:
+            query = self.switch(self.model_class)
+        else:
+            query = self.clone()
         query = query.ensure_join(query._query_ctx, rel_model)
         if not query._group_by:
             query._group_by = [x.alias() for x in query._select]

diff --git a/playhouse/tests/test_queries.py b/playhouse/tests/test_queries.py
@@ -1045,6 +1045,12 @@ def test_annotate(self):
         self.assertWhere(sq, '', [])
         self.assertGroupBy(sq, '"users"."id", "users"."username"', [])
 
+        sq = User.select().join(Blog, JOIN_LEFT_OUTER).annotate(Blog)
+        self.assertSelect(sq, '"users"."id", "users"."username", Count("blog"."pk") AS count', [])
+        self.assertJoins(sq, ['LEFT OUTER JOIN "blog" AS blog ON ("users"."id" = "blog"."user_id")'])
+        self.assertWhere(sq, '', [])
+        self.assertGroupBy(sq, '"users"."id", "users"."username"', [])
+
     def test_aggregate(self):
         sq = User.select().where(User.id < 10)._aggregate()
         self.assertSelect(sq, 'Count(*)', [])