Add count methods to context and query

This includes tests for these methods.

Signed-off-by: Don Naegely <naegelyd@gmail.com>

avocado/query/models.py

@@ -1,6 +1,7 @@
 import jsonfield
 from django.db import models
 from modeltree.tree import trees
+from avocado.core.cache import cached_method
 from . import oldparsers as parsers
 
 
@@ -68,6 +69,10 @@ def validate(cls, attrs, **context):
         "Validate `attrs` as a context."
         return parsers.datacontext.validate(attrs, **context)
 
+    @cached_method(version='modified')
+    def count(self, *args, **kwargs):
+        return self.apply(*args, **kwargs).count()
+
     def parse(self, tree=None, **context):
         "Returns a parsed node for this context."
         return parsers.datacontext.parse(self.json, tree=tree, **context)
@@ -199,6 +204,10 @@ def validate(cls, attrs, **context):
         "Validates `attrs` as a query."
         return parsers.dataquery.validate(attrs, **context)
 
+    @cached_method(version='modified')
+    def count(self, *args, **kwargs):
+        return self.apply(*args, **kwargs).count()
+
     def parse(self, tree=None, **context):
         "Returns a parsed node for this query."
         json = {
@@ -212,6 +221,7 @@ def apply(self, queryset=None, tree=None, distinct=True, include_pk=True,
         "Applies this context to a QuerySet."
         if tree is None and queryset is not None:
             tree = queryset.model
+
         return self.parse(tree=tree, **context) \
             .apply(queryset=queryset, distinct=distinct, include_pk=include_pk)
 

tests/cases/models/tests.py

@@ -464,6 +464,11 @@ def test_published(self):
 
 
 class DataContextTestCase(TestCase):
+    fixtures = ['employee_data.json']
+
+    def setUp(self):
+        management.call_command('avocado', 'init', 'tests')
+
     def test_init(self):
         json = {
             'field': 'tests.title.salary',
@@ -488,6 +493,17 @@ def test_clean(self):
 
         cxt.save()
 
+    def test_count(self):
+        json = {
+            'field': 'tests.title.salary',
+            'operator': 'gt',
+            'value': '10000'
+        }
+        ctx = DataContext(json)
+
+        self.assertEqual(ctx.count(), 6)
+        self.assertEqual(ctx.count(tree='office'), 1)
+
 
 class DataViewTestCase(TestCase):
     def test_init(self):
@@ -512,23 +528,16 @@ def test_clean(self):
 
 
 class DataQueryTestCase(TestCase):
+    fixtures = ['employee_data.json']
+
     existing_email = 'existing@email.com'
     existing_username = 'user1'
     emails = [existing_email, 'new1@email.com', 'new2@email.com',
               'new3@email.com']
     usernames = [existing_username, 'user3', 'user4', 'user5', 'user6']
 
     def setUp(self):
-        management.call_command('avocado', 'init', 'tests', publish=False,
-                                concepts=False, quiet=True)
-        f1 = DataField.objects.get(pk=1)
-        f2 = DataField.objects.get(pk=2)
-
-        c1 = DataConcept()
-        c1.save()
-
-        DataConceptField(concept=c1, field=f1).save()
-        DataConceptField(concept=c1, field=f2).save()
+        management.call_command('avocado', 'init', 'tests')
 
     def test_init(self):
         json = {
@@ -550,6 +559,34 @@ def test_init(self):
 
         self.assertEqual(query.json, json)
 
+    def test_count(self):
+        salary_field = DataField.objects.get_by_natural_key(
+            'tests', 'title', 'salary')
+
+        json = {
+            'context': {
+                'field': 'tests.title.salary',
+                'operator': 'gt',
+                'value': '1000'
+            },
+            'view': {'columns': [salary_field.id]}
+        }
+
+        query = DataQuery(json)
+
+        # Default tree is Employee so we should get 6 unique employee objects
+        # regardless of distinct setting since all are distinct.
+        self.assertEqual(query.count(), 6)
+        self.assertEqual(query.count(distinct=False), 6)
+
+        # Switching the tree should allow us to exercise the distinct keyword
+        # since there are 3 titles all with the same 15,000 unit salary. We
+        # need to eliminate the PK so that we are only getting salaries
+        # back. Including the PK causes everything to be unique.
+        self.assertEqual(query.count(tree='title', include_pk=False), 5)
+        self.assertEqual(
+            query.count(tree='title', include_pk=False, distinct=False), 7)
+
     def test_multiple_json_values(self):
         json = {
             'context': {
@@ -630,7 +667,7 @@ def test_apply(self):
         self.assertEqual(
             unicode(query.apply(tree=Employee).query),
             'SELECT DISTINCT "tests_employee"."id", '
-            '"tests_office"."location", "tests_title"."name" FROM '
+            '"tests_office"."location" FROM '
             '"tests_employee" INNER JOIN "tests_title" ON '
             '("tests_employee"."title_id" = "tests_title"."id") INNER JOIN '
             '"tests_office" ON ("tests_employee"."office_id" = '
@@ -641,12 +678,10 @@ def test_apply(self):
         self.assertEqual(
             unicode(query.apply(queryset=queryset).query),
             'SELECT DISTINCT "tests_employee"."id", '
-            '"tests_office"."location", "tests_title"."name" FROM '
+            '"tests_office"."location" FROM '
             '"tests_employee" INNER JOIN "tests_office" ON '
-            '("tests_employee"."office_id" = "tests_office"."id") LEFT OUTER '
-            'JOIN "tests_title" ON ("tests_employee"."title_id" = '
-            '"tests_title"."id") ORDER BY "tests_office"."location" DESC, '
-            '"tests_title"."name" DESC')
+            '("tests_employee"."office_id" = "tests_office"."id") '
+            'ORDER BY "tests_office"."location" DESC')
 
     def test_clean(self):
         # Save default template

tests/settings.py

@@ -88,6 +88,12 @@
 MODELTREES = {
     'default': {
         'model': 'tests.Employee',
+    },
+    'title': {
+        'model': 'tests.Title',
+    },
+    'office': {
+        'model': 'tests.Office',
     }
 }