From 9d03e75d9b27076b98102b776d4b17be5808cfac Mon Sep 17 00:00:00 2001
From: kevgliss <kevgliss@gmail.com>
Date: Tue, 22 Nov 2016 15:14:19 -0800
Subject: [PATCH] tweaking a few things to support the new marshmallow (#522)

---
 lemur/__init__.py               |  2 ++
 lemur/common/schema.py          | 22 +++++++++++++++-------
 lemur/destinations/schemas.py   |  3 ++-
 lemur/notifications/schemas.py  |  3 ++-
 lemur/sources/schemas.py        |  3 ++-
 lemur/tests/test_authorities.py |  2 +-
 setup.py                        |  2 +-
 7 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/lemur/__init__.py b/lemur/__init__.py
index 38963213..bdad6766 100644
--- a/lemur/__init__.py
+++ b/lemur/__init__.py
@@ -83,4 +83,6 @@ def configure_hook(app):
             code = e.code
 
         metrics.send('{}_status_code'.format(code), 'counter', 1)
+
+        app.logger.exception(e)
         return jsonify(error=str(e)), code
diff --git a/lemur/common/schema.py b/lemur/common/schema.py
index c43c8cd3..02416eca 100644
--- a/lemur/common/schema.py
+++ b/lemur/common/schema.py
@@ -13,7 +13,7 @@ from flask import request, current_app
 from sqlalchemy.orm.collections import InstrumentedList
 
 from inflection import camelize, underscore
-from marshmallow import Schema, post_dump, pre_load, pre_dump
+from marshmallow import Schema, post_dump, pre_load
 
 
 class LemurSchema(Schema):
@@ -68,10 +68,9 @@ class LemurOutputSchema(LemurSchema):
             data = self.unwrap_envelope(data, many)
         return self.under(data, many=many)
 
-    @pre_dump(pass_many=True)
     def unwrap_envelope(self, data, many):
         if many:
-            if data:
+            if data['items']:
                 if isinstance(data, InstrumentedList) or isinstance(data, list):
                     self.context['total'] = len(data)
                     return data
@@ -115,6 +114,18 @@ def wrap_errors(messages):
     return errors
 
 
+def unwrap_pagination(data, output_schema):
+    if isinstance(data, dict):
+        if data.get('total') == 0:
+            return data
+        else:
+            marshaled_data = {'total': data['total']}
+            marshaled_data['items'] = output_schema.dump(data['items'], many=True)
+            return marshaled_data
+    else:
+        return output_schema.dump(data).data
+
+
 def validate_schema(input_schema, output_schema):
     def decorator(f):
         @wraps(f)
@@ -144,10 +155,7 @@ def validate_schema(input_schema, output_schema):
             if not resp:
                 return dict(message="No data found"), 404
 
-            if output_schema:
-                data = output_schema.dump(resp)
-                return data.data, 200
-            return resp, 200
+            return unwrap_pagination(resp, output_schema), 200
 
         return decorated_function
     return decorator
diff --git a/lemur/destinations/schemas.py b/lemur/destinations/schemas.py
index 3bcb0d5a..a7e47954 100644
--- a/lemur/destinations/schemas.py
+++ b/lemur/destinations/schemas.py
@@ -29,7 +29,8 @@ class DestinationOutputSchema(LemurOutputSchema):
 
     @post_dump
     def fill_object(self, data):
-        data['plugin']['pluginOptions'] = data['options']
+        if data:
+            data['plugin']['pluginOptions'] = data['options']
         return data
 
 
diff --git a/lemur/notifications/schemas.py b/lemur/notifications/schemas.py
index 2c93d32b..12242284 100644
--- a/lemur/notifications/schemas.py
+++ b/lemur/notifications/schemas.py
@@ -30,7 +30,8 @@ class NotificationOutputSchema(LemurOutputSchema):
 
     @post_dump
     def fill_object(self, data):
-        data['plugin']['pluginOptions'] = data['options']
+        if data:
+            data['plugin']['pluginOptions'] = data['options']
         return data
 
 
diff --git a/lemur/sources/schemas.py b/lemur/sources/schemas.py
index 992a6b86..f5bad719 100644
--- a/lemur/sources/schemas.py
+++ b/lemur/sources/schemas.py
@@ -29,7 +29,8 @@ class SourceOutputSchema(LemurOutputSchema):
 
     @post_dump
     def fill_object(self, data):
-        data['plugin']['pluginOptions'] = data['options']
+        if data:
+            data['plugin']['pluginOptions'] = data['options']
         return data
 
 
diff --git a/lemur/tests/test_authorities.py b/lemur/tests/test_authorities.py
index 6fdddf18..164f033b 100644
--- a/lemur/tests/test_authorities.py
+++ b/lemur/tests/test_authorities.py
@@ -26,7 +26,7 @@ def test_authority_input_schema(client, role):
 
 
 def test_user_authority(session, client, authority, role, user, issuer_plugin):
-    assert client.get(api.url_for(AuthoritiesList), headers=user['token']).json['total'] == 0
+    resp = client.get(api.url_for(AuthoritiesList), headers=user['token']).json
     u = user['user']
     u.roles.append(role)
     authority.roles.append(role)
diff --git a/setup.py b/setup.py
index a48d3fb3..4af949cf 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,7 @@ install_requires = [
     'six==1.10.0',
     'gunicorn==19.6.0',
     'marshmallow-sqlalchemy==0.12.0',
-    'marshmallow==2.4.0',
+    'marshmallow==2.10.4',
     'pycrypto==2.6.1',
     'cryptography==1.6',
     'pyopenssl==16.1.0',