Handling type row values

Author: cmanallen

jsonapi_query/database/sqlalchemy.py

@@ -183,6 +183,8 @@ def group_and_remove(items, models):
     response = [[] for model in models]
     for item in items:
         for member in item:
+            if member is None:
+                continue
             position = models.index(member.__class__)
             if member not in response[position]:
                 response[position].append(member)

tests/unit/database/sqlalchemy_tests.py

@@ -3,7 +3,7 @@
 
 from sqlalchemy.orm import Query, sessionmaker
 
-from jsonapi_query.database.sqlalchemy import QueryMixin
+from jsonapi_query.database.sqlalchemy import group_and_remove, QueryMixin
 from tests.sqlalchemy import (
     BaseSQLAlchemyTestCase, Category, Person, Product, School, Student)
 
@@ -287,14 +287,47 @@ def test_include_ambiguous_join_conditions(self):
 
     def test_include_does_not_restrict_primary_output(self):
         """Test including a relationship does not restrict primary output."""
-        p = Product(name='Tst')
+        a = Category(name='Category A')
+        self.session.add(a)
+        b = Category(name='Category B', category_id=1)
+        self.session.add(b)
+        p = Product(primary_category_id=1, secondary_category_id=2, name='Tst')
+        self.session.add(p)
+
+        p = Product(name='Tst 2')
         self.session.add(p)
 
         models = self.session.query(Product).include(
             [Product.primary_category]).all()
-        self.assertTrue(len(models) == 1)
+        self.assertTrue(len(models) == 2)
 
     def test_include_no_mappers(self):
         """Test including an empty set of relationships."""
         models = self.session.query(Person).include([]).first()
         self.assertTrue(isinstance(models, Person))
+
+
+class UtilitySQLAlchemyTestCase(BaseDatabaseSQLAlchemyTests):
+    """Test handling a query's output."""
+
+    def test_group_and_remove(self):
+        """Test group and remove utility function."""
+        a = Category(name='Category A')
+        self.session.add(a)
+        b = Category(name='Category B', category_id=1)
+        self.session.add(b)
+        p = Product(primary_category_id=1, secondary_category_id=2, name='Tst')
+        self.session.add(p)
+
+        p = Product(name='Tst 2')
+        self.session.add(p)
+
+        # Returns two products and one category.
+        items = self.session.query(Product).include(
+            [Product.primary_category]).all()
+        self.assertTrue(len(items) == 2)
+
+        output = group_and_remove(items, [Product, Category])
+        self.assertTrue(len(output) == 2)
+        self.assertTrue(len(output[0]) == 2)
+        self.assertTrue(len(output[1]) == 1)