Card #30 - migrates attachments controller to twisted

3 files changed, 69 insertions, 8 deletions

diff --git a/service/pixelated/controllers/attachments_controller.py b/service/pixelated/controllers/attachments_controller.py
index 1d5360f7..4965f166 100644
--- a/service/pixelated/controllers/attachments_controller.py
+++ b/service/pixelated/controllers/attachments_controller.py
@@ -14,12 +14,11 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
 
-from flask import send_file
-from flask import request
-
 import io
 import re
-from twisted.web.server import NOT_DONE_YET
+
+from twisted.protocols.basic import FileSender
+from twisted.python.log import err
 
 
 class AttachmentsController:
@@ -28,12 +27,21 @@ class AttachmentsController:
         self.querier = querier
 
     def attachment(self, request, attachment_id):
-        encoding = request.args.get('encoding', [''])[0]
+        encoding = request.args.get('encoding')
         attachment = self.querier.attachment(attachment_id, encoding)
-        request.setRawHeader('Content-Type', self._extract_mimetype(attachment['content-type']))
-        request.write(io.BytesIO(attachment['content']))
 
-        return NOT_DONE_YET
+        content_type = self._extract_mimetype(attachment['content-type'])
+        request.setHeader('Content-Type', content_type)
+        bytes_io = io.BytesIO(attachment['content'])
+        d = FileSender().beginFileTransfer(bytes_io, request)
+
+        def cbFinished(ignored):
+            bytes_io.close()
+            request.finish()
+
+        d.addErrback(err).addCallback(cbFinished)
+
+        return d
 
     def _extract_mimetype(self, content_type):
         match = re.compile('([A-Za-z-]+\/[A-Za-z-]+)').search(content_type)
diff --git a/service/test/integration/retrieve_attachment_test.py b/service/test/integration/retrieve_attachment_test.py
new file mode 100644
index 00000000..37bc3ca2
--- /dev/null
+++ b/service/test/integration/retrieve_attachment_test.py
@@ -0,0 +1,42 @@
+#
+# Copyright (c) 2014 ThoughtWorks, Inc.
+#
+# Pixelated is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pixelated is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
+import unittest
+
+from test.support.integration_helper import SoledadTestBase
+
+
+class RetrieveAttachmentTest(unittest.TestCase, SoledadTestBase):
+
+    def setUp(self):
+        self.setup_soledad()
+
+    def tearDown(self):
+        self.teardown_soledad()
+
+    def test_attachment_content_is_retrieved(self):
+        ident = 'F4E99C1CEC4D300A4223A96CCABBE0304BDBC31C550A5A03E207A5E4C3C71A22'
+        attachment_dict = {'content-disposition': 'attachment',
+                           'content-transfer-encoding': '',
+                           'type': 'cnt',
+                           'raw': 'cGVxdWVubyBhbmV4byA6RAo=',
+                           'phash': ident,
+                           'content-type': 'text/plain; charset=US-ASCII; name="attachment_pequeno.txt"'}
+
+        self.add_document_to_soledad(attachment_dict)
+
+        attachment = self.get_attachment(ident, 'base64')
+
+        self.assertEquals('pequeno anexo :D\n', attachment)
diff --git a/service/test/support/integration_helper.py b/service/test/support/integration_helper.py
index 6ab96e52..36d297d6 100644
--- a/service/test/support/integration_helper.py
+++ b/service/test/support/integration_helper.py
@@ -183,6 +183,14 @@ class SoledadTestBase:
         response = json.loads(request.getWrittenData())
         return ResponseMail(response)
 
+    def get_attachment(self, ident, encoding):
+        request = requestMock(path='/attachment/' + ident)
+        request.args = {
+            'encoding': encoding
+        }
+        _render(self.resource, request)
+        return request.getWrittenData()
+
     def put_mail(self, data):
         request = requestMock('/mails', method="PUT", body=data, headers={'Content-Type': ['application/json']})
         _render(self.resource, request)
@@ -226,6 +234,9 @@ class SoledadTestBase:
         _render(self.resource, request)
         return request
 
+    def add_document_to_soledad(self, _dict):
+        self.soledad_querier.soledad.create_doc(_dict)
+
     def add_mail_to_inbox(self, input_mail):
         mail = self.mailboxes.inbox().add(input_mail)
         mail.update_tags(input_mail.tags)