Create new Attachment.content_disposition

loadsmart · Jan 4, 2024 · 116651b · 116651b
1 parent 27061e4
commit 116651b
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 5 deletions.
diff --git a/changelogs/unreleased/fops-5403.yml b/changelogs/unreleased/fops-5403.yml
@@ -0,0 +1,4 @@
+---
+description: 'The killerz: Create new Attachment.content_disposition'
+type: added
+pr_number: 30
diff --git a/gmail_wrapper/entities.py b/gmail_wrapper/entities.py
@@ -1,4 +1,5 @@
 import base64
+import cgi
 from datetime import datetime
 
 
@@ -29,6 +30,7 @@ def __init__(self, message_id, client, raw_part):
         self._raw = raw_part
         self._client = client
         self._body = AttachmentBody(raw_part["body"])
+        self._headers = raw_part.get("headers", [])
         self.message_id = message_id
 
     @property
@@ -50,6 +52,25 @@ def content(self):
 
         return self._body.content
 
+    @property
+    def content_disposition(self):
+        content_disposition_header = next(
+            (
+                header
+                for header in self._headers
+                if header.get("name") == "Content-Disposition"
+            ),
+            {},
+        )
+
+        content_disposition_value = content_disposition_header.get("value")
+
+        return (
+            cgi.parse_header(content_disposition_value)[0]
+            if content_disposition_value
+            else None
+        )
+
 
 class Message:
     def __init__(self, client, raw_message):

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -47,7 +47,10 @@ def raw_complete_message():
                 {"name": "To", "value": "foo@loadsmart.com"},
                 {"name": "From", "value": "john@doe.com"},
                 {"name": "Subject", "value": "Urgent errand"},
-                {"name": "Message-ID", "value": "<BY5PR15MB353717D866FC27FEE4DB4EC7F77E0@BY5PR15MB3537.namprd15.prod.outlook.com>"},
+                {
+                    "name": "Message-ID",
+                    "value": "<BY5PR15MB353717D866FC27FEE4DB4EC7F77E0@BY5PR15MB3537.namprd15.prod.outlook.com>",
+                },
             ],
             "parts": [
                 {
@@ -64,6 +67,12 @@ def raw_complete_message():
                     "body": {"data": "", "attachmentId": "CCX457", "size": 60},
                     "mimeType": "text/plain",
                     "partId": "BB790",
+                    "headers": [
+                        {
+                            "name": "Content-Disposition",
+                            "value": "inline",
+                        },
+                    ],
                     "filename": "fox.txt",
                 },
                 {
@@ -81,8 +90,26 @@ def raw_complete_message():
                             "body": {"attachmentId": "CCX458", "size": 95},
                             "mimeType": "application/pdf",
                             "partId": "BB791.1",
+                            "headers": [
+                                {
+                                    "name": "Content-Disposition",
+                                    "value": 'attachment; filename="tigers.pdf"; size=95',
+                                },
+                            ],
                             "filename": "tigers.pdf",
-                        }
+                        },
+                        {
+                            "partId": "BB791.2",
+                            "mimeType": "image/jpeg",
+                            "filename": "image001.jpg",
+                            "headers": [
+                                {
+                                    "name": "Content-Disposition",
+                                    "value": 'inline; filename="image001.jpg"; size=48',
+                                },
+                            ],
+                            "body": {"attachmentId": "ANGjdJ", "size": 48},
+                        },
                     ],
                 },
             ],

diff --git a/tests/test_entities.py b/tests/test_entities.py
@@ -77,11 +77,16 @@ def test_it_does_not_generate_attachment_objects_for_attachments_without_id(
                 "mimeType": "application/pdf",
             }
         )
-        assert len(complete_message.attachments) == 2
-        assert complete_message.attachments[0].id
+        assert len(complete_message.attachments) == 3
+        assert complete_message.attachments[0].id == "CCX457"
         assert complete_message.attachments[0].filename == "fox.txt"
-        assert complete_message.attachments[1].id
+        assert complete_message.attachments[0].content_disposition == "inline"
+        assert complete_message.attachments[1].id == "CCX458"
         assert complete_message.attachments[1].filename == "tigers.pdf"
+        assert complete_message.attachments[1].content_disposition == "attachment"
+        assert complete_message.attachments[2].id == "ANGjdJ"
+        assert complete_message.attachments[2].filename == "image001.jpg"
+        assert complete_message.attachments[2].content_disposition == "inline"
 
     def test_it_returns_empty_list_if_no_attachments(
         self, client, raw_complete_message
@@ -204,6 +209,9 @@ def test_it_fetch_additional_information_when_needed(
             "123AAB", client, raw_complete_message["payload"]["parts"][1]
         )
         assert incomplete_attachment.content
+        assert incomplete_attachment.id == "CCX457"
+        assert incomplete_attachment.filename == "fox.txt"
+        assert incomplete_attachment.content_disposition == "inline"
         mocked_get_attachment_body.assert_called_once_with(
             raw_attachment_body["attachmentId"], "123AAB"
         )