Merge pull request #106 from cloudant/102-fix-get-attachment-bug

Merged

CHANGES.rst

@@ -4,6 +4,7 @@
 - [BREAKING] Fixed CloudantDatabase.share_database to accept all valid permission roles.  Changed the method signature to accept roles as a list argument.
 - [FIX] The CouchDatabase.create_document method now will handle both documents and design documents correctly.  If the document created is a design document then the locally cached object will be a DesignDocument otherwise it will be a Document.
 - [BREAKING] Refactored the SearchIndex class to now be the TextIndex class.  Also renamed the CloudantDatabase convenience methods of get_all_indexes, create_index, and delete_index as get_query_indexes, create_query_index, and delete_query_index respectively.  These changes were made to clarify that the changed class and the changed methods were specific to query index processing only.
+- [FIX] Fixed Document.get_attachment method to successfully create text and binary files based on http response Content-Type.  The method also returns text, binary, and json content based on http response Content-Type.
 
 2.0.0b2 (2016-02-24)
 ====================

src/cloudant/document.py

@@ -20,7 +20,7 @@
 import requests
 from requests.exceptions import HTTPError
 
-from ._2to3 import unicode_, url_quote, url_quote_plus
+from ._2to3 import url_quote, url_quote_plus
 from .errors import CloudantException
 
 
@@ -363,9 +363,12 @@ def get_attachment(
             attachment,
             headers=None,
             write_to=None,
-            attachment_type="json"):
+            attachment_type=None):
         """
         Retrieves a document's attachment and optionally writes it to a file.
+        If the content_type of the attachment is 'application/json' then the
+        data returned will be in JSON format otherwise the response content will
+        be returned as text or binary.
 
         :param str attachment: Attachment file name used to identify the
             attachment.
@@ -374,8 +377,11 @@ def get_attachment(
         :param file write_to: Optional file handler to write the attachment to.
             The write_to file must be opened for writing prior to including it
             as an argument for this method.
-        :param str attachment_type: Data format of the attachment.  Valid
-            values are ``'json'`` and ``'binary'``.
+        :param str attachment_type: Optional setting to define how to handle the
+            attachment when returning its contents from this method.  Valid
+            values are ``'text'``, ``'json'``, and ``'binary'``  If
+            omitted then the returned content will be based on the
+            response Content-Type.
 
         :returns: The attachment content
         """
@@ -387,17 +393,28 @@ def get_attachment(
         else:
             headers['If-Match'] = self['_rev']
 
-        resp = self.r_session.get(
-            attachment_url,
-            headers=headers
-        )
+        resp = self.r_session.get(attachment_url, headers=headers)
         resp.raise_for_status()
-        if write_to is not None:
-            write_to.write(resp.raw)
 
-        if attachment_type == 'json':
+        if attachment_type is None:
+            if resp.headers['Content-Type'].startswith('text/'):
+                attachment_type = 'text'
+            elif resp.headers['Content-Type'] == 'application/json':
+                attachment_type = 'json'
+            else:
+                attachment_type = 'binary'
+
+        if write_to is not None:
+            if attachment_type == 'text' or attachment_type == 'json':
+                write_to.write(resp.text)
+            else:
+                write_to.write(resp.content)
+        if attachment_type == 'text':
+            return resp.text
+        elif attachment_type == 'json':
             return resp.json()
-        return unicode_(resp.content)
+
+        return resp.content
 
     def delete_attachment(self, attachment, headers=None):
         """

tests/unit/db/document_tests.py

@@ -26,13 +26,22 @@
 import posixpath
 import json
 import requests
+import os
+import uuid
+import inspect
 
 from cloudant.document import Document
 from cloudant.errors import CloudantException
 
-from ... import StringIO
+from ... import StringIO, unicode_
 from .unit_t_db_base import UnitTestDbBase
 
+def find_fixture(name):
+    import tests.unit.db.fixtures as fixtures
+    dirname = os.path.dirname(inspect.getsourcefile(fixtures))
+    filename = os.path.join(dirname, name)
+    return filename
+
 class DocumentTests(UnitTestDbBase):
     """
     Document unit tests
@@ -469,6 +478,77 @@ def test_removing_id(self):
         self.assertIsNone(doc.get('_id'))
         self.assertEqual(doc._document_id, None)
 
+    def test_get_text_attachment(self):
+        """
+        Test the retrieval of a text attachment
+        """
+        doc = self.db.create_document(
+            {'_id': 'julia006', 'name': 'julia', 'age': 6}
+        )
+        attachment = StringIO()
+        try:
+            filename = 'attachment-{0}{1}'.format(unicode_(uuid.uuid4()), '.txt')
+            attachment.write('This is line one of the attachment.\n')
+            attachment.write('This is line two of the attachment.\n')
+            resp = doc.put_attachment(
+                filename,
+                'text/plain',
+                attachment.getvalue()
+            )
+            with open(find_fixture(filename), 'wt') as f:
+                text_attachment = doc.get_attachment(filename, write_to=f)
+                self.assertEqual(text_attachment, attachment.getvalue())
+            with open(find_fixture(filename), 'rt') as f:
+                self.assertEqual(f.read(), attachment.getvalue())
+        finally:
+            attachment.close()
+            os.remove(find_fixture(filename))
+
+    def test_get_json_attachment(self):
+        """
+        Test the retrieval of a json attachment
+        """
+        doc = self.db.create_document(
+            {'_id': 'julia006', 'name': 'julia', 'age': 6}
+        )
+        try:
+            filename = 'attachment-{0}{1}'.format(unicode_(uuid.uuid4()), '.json')
+            data = {'foo': 'bar', 'baz': 99}
+            resp = doc.put_attachment(
+                filename,
+                'application/json',
+                json.dumps(data)
+            )
+            with open(find_fixture(filename), 'wt') as f:
+                json_attachment = doc.get_attachment(filename, write_to=f)
+                self.assertIsInstance(json_attachment, dict)
+                self.assertEqual(json_attachment, data)
+            with open(find_fixture(filename), 'rt') as f:
+                self.assertEqual(f.read(), json.dumps(data))
+        finally:
+            os.remove(find_fixture(filename))
+
+    def test_get_binary_attachment(self):
+        """
+        Test the retrieval of a binary attachment
+        """
+        doc = self.db.create_document(
+            {'_id': 'julia006', 'name': 'julia', 'age': 6}
+        )
+        try:
+            filename = 'attachment-{0}{1}'.format(unicode_(uuid.uuid4()), '.jpg')
+            data = None
+            with open(find_fixture('smile.jpg'), 'rb') as f:
+                data = f.read()
+                resp = doc.put_attachment(filename,'image/jpeg', data)
+            with open(find_fixture(filename), 'wb') as f:
+                binary_attachment = doc.get_attachment(filename, write_to=f)
+                self.assertEqual(binary_attachment, data)
+            with open(find_fixture(filename), 'rb') as f:
+                self.assertEqual(f.read(), data)
+        finally:
+            os.remove(find_fixture(filename))
+
     def test_attachment_management(self):
         """
         Test the adding, retrieving, updating, and deleting of attachments

tests/unit/db/fixtures/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+# Copyright (c) 2015 IBM. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); 
+# you may not use this file except in compliance with the License. 
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software 
+# distributed under the License is distributed on an "AS IS" BASIS, 
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+# See the License for the specific language governing permissions and 
+# limitations under the License.
+"""
+_fixtures_
+
+Fixtures folder containing files used by db tests
+
+"""