diff options
-rw-r--r-- | service/pixelated/controllers/attachments_controller.py | 24 | ||||
-rw-r--r-- | service/test/integration/retrieve_attachment_test.py | 42 | ||||
-rw-r--r-- | service/test/support/integration_helper.py | 11 |
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) |