Merge pull request #10 from loadsmart/filterAttachmentsWithoutId

Filter attachments without ID

Author: luizguilhermefr

gmail_wrapper/entities.py

@@ -112,7 +112,7 @@ def attachments(self):
         return [
             Attachment(self.id, self._client, part)
             for part in (parts if parts else [])
-            if part["filename"]
+            if part["filename"] and part["body"] and part["body"].get("attachmentId")
         ]
 
     def modify(self, add_labels=None, remove_labels=None):

tests/test_entities.py

@@ -65,6 +65,22 @@ def test_it_generates_attachment_objects(self, client, raw_complete_message):
             ]
         )
 
+    def test_it_does_not_generate_attachment_objects_for_attachments_without_id(
+        self, client, raw_complete_message
+    ):
+        complete_message = Message(client, raw_complete_message)
+        raw_complete_message["payload"]["parts"].append(
+            {
+                "body": {"size": 0, "data": ""},
+                "filename": "invalid.pdf",
+                "partId": "BB710",
+                "mimeType": "application/pdf",
+            }
+        )
+        assert len(complete_message.attachments) == 1
+        assert complete_message.attachments[0].id
+        assert complete_message.attachments[0].filename != "invalid.pdf"
+
     def test_it_returns_empty_list_if_no_attachments(
         self, client, raw_complete_message
     ):