Merge pull request #8 from loadsmart/handleExceptions

Handle Gmail exceptions

Author: luizguilhermefr

README.md

@@ -104,6 +104,27 @@ message = client.send(
 print(message) # Gmail message: ABC123
 ```
 
+- Handle exceptions
+
+Exceptions are part of every developer day-to-day. You may want to handle exceptions as follows:
+
+```python
+from gmail_wrapper.exceptions import (
+    MessageNotFoundError,
+    AttachmentNotFoundError,
+    GmailError,
+)
+
+try:
+    do_something()
+except MessageNotFoundError as e:
+    print(f"There is no message! {e}")
+except AttachmentNotFoundError as e:
+    print(f"There is no attachment! {e}")
+except GmailError as e:
+    print(f"Google servers are burning! {e}")
+```
+
 ## Need more help?
 
 Reach `#eng-accounting` Slack channel.

gmail_wrapper/client.py

@@ -5,8 +5,14 @@
 from google.auth.transport.requests import Request
 from google.oauth2.credentials import Credentials
 from googleapiclient import discovery
+from googleapiclient.errors import HttpError
 
 from gmail_wrapper.entities import Message, AttachmentBody
+from gmail_wrapper.exceptions import (
+    MessageNotFoundError,
+    AttachmentNotFoundError,
+    GmailError,
+)
 
 
 class GmailClient:
@@ -38,11 +44,17 @@ def _make_client(self, secrets_json_string, scopes):
     def _messages_resource(self):
         return self._client.users().messages()
 
+    def _execute(self, executable):
+        try:
+            return executable.execute()
+        except HttpError as e:
+            if e.resp.status >= 500:
+                raise GmailError()
+            raise e
+
     def get_raw_messages(self, query="", limit=None):
-        return (
-            self._messages_resource()
-            .list(userId=self.email, q=query, maxResults=limit)
-            .execute()
+        return self._execute(
+            self._messages_resource().list(userId=self.email, q=query, maxResults=limit)
         )
 
     def get_messages(self, query="", limit=None):
@@ -54,49 +66,65 @@ def get_messages(self, query="", limit=None):
         return [Message(self, raw_message) for raw_message in raw_messages["messages"]]
 
     def get_raw_message(self, id):
-        return self._messages_resource().get(userId=self.email, id=id).execute()
+        try:
+            return self._execute(
+                self._messages_resource().get(userId=self.email, id=id)
+            )
+        except HttpError as e:
+            if e.resp.status == 404:
+                raise MessageNotFoundError(id)
+            raise e
 
     def get_message(self, id):
         raw_message = self.get_raw_message(id)
 
         return Message(self, raw_message)
 
     def modify_raw_message(self, id, add_labels=None, remove_labels=None):
-        return (
-            self._messages_resource()
-            .modify(
-                userId=self.email,
-                id=id,
-                body={
-                    "addLabelIds": add_labels if add_labels else [],
-                    "removeLabelIds": remove_labels if remove_labels else [],
-                },
+        try:
+            return self._execute(
+                self._messages_resource().modify(
+                    userId=self.email,
+                    id=id,
+                    body={
+                        "addLabelIds": add_labels if add_labels else [],
+                        "removeLabelIds": remove_labels if remove_labels else [],
+                    },
+                )
             )
-            .execute()
-        )
+        except HttpError as e:
+            if e.resp.status == 404:
+                raise MessageNotFoundError(id)
+            raise e
 
     def modify_message(self, id, add_labels=None, remove_labels=None):
         raw_modified_message = self.modify_raw_message(id, add_labels, remove_labels)
 
         return Message(self, raw_modified_message)
 
     def modify_multiple_messages(self, ids, add_labels=None, remove_labels=None):
-        self._messages_resource().batchModify(
-            userId=self.email,
-            body={
-                "ids": ids,
-                "addLabelIds": add_labels if add_labels else [],
-                "removeLabelIds": remove_labels if remove_labels else [],
-            },
-        ).execute()
+        self._execute(
+            self._messages_resource().batchModify(
+                userId=self.email,
+                body={
+                    "ids": ids,
+                    "addLabelIds": add_labels if add_labels else [],
+                    "removeLabelIds": remove_labels if remove_labels else [],
+                },
+            )
+        )
 
     def get_raw_attachment_body(self, id, message_id):
-        return (
-            self._messages_resource()
-            .attachments()
-            .get(userId=self.email, id=id, messageId=message_id)
-            .execute()
-        )
+        try:
+            return self._execute(
+                self._messages_resource()
+                .attachments()
+                .get(userId=self.email, id=id, messageId=message_id)
+            )
+        except HttpError as e:
+            if e.resp.status == 404:
+                raise AttachmentNotFoundError(message_id, id)
+            raise e
 
     def get_attachment_body(self, id, message_id):
         raw_attachment_body = self.get_raw_attachment_body(id, message_id)
@@ -121,8 +149,8 @@ def send_raw(self, subject, html_content, to, cc=None, bcc=None):
             subject, html_content, to, cc if cc else [], bcc if bcc else []
         )
 
-        return (
-            self._messages_resource().send(userId=self.email, body=sendable).execute()
+        return self._execute(
+            self._messages_resource().send(userId=self.email, body=sendable)
         )
 
     def send(self, subject, html_content, to, cc=None, bcc=None):

gmail_wrapper/exceptions.py

@@ -0,0 +1,20 @@
+class GmailError(Exception):
+    def __str__(self):
+        return f"GmailError: Gmail returned an internal server error"
+
+
+class MessageNotFoundError(Exception):
+    def __init__(self, message_id):
+        self.message_id = message_id
+
+    def __str__(self):
+        return f"MessageNotFoundError: Gmail returned 404 when attempting to get message {self.message_id}"
+
+
+class AttachmentNotFoundError(Exception):
+    def __init__(self, message_id, attachment_id):
+        self.message_id = message_id
+        self.attachment_id = attachment_id
+
+    def __str__(self):
+        return f"AttachmentNotFoundError: Gmail returned 404 when attempting to get attachment {self.attachment_id} of message {self.message_id}"

tests/test_client.py

@@ -1,7 +1,15 @@
 import base64
 
+import pytest
+from googleapiclient.errors import HttpError
+
 from gmail_wrapper import GmailClient
 from gmail_wrapper.entities import Message, AttachmentBody
+from gmail_wrapper.exceptions import (
+    MessageNotFoundError,
+    AttachmentNotFoundError,
+    GmailError,
+)
 from tests.utils import make_gmail_client
 
 
@@ -15,6 +23,19 @@ def test_it_returns_raw_messages(self, mocker, raw_incomplete_message):
         client = GmailClient(email="[email protected]", secrets_json_string="{}")
         assert client.get_raw_messages() == raw_response
 
+    def test_it_encapsulates_gmail_exceptions(self, mocker):
+        server_error_response = mocker.MagicMock(status=500)
+        mocker.patch(
+            "gmail_wrapper.client.GmailClient._make_client",
+            return_value=make_gmail_client(
+                mocker, list_effect=HttpError(server_error_response, b"Content")
+            ),
+        )
+
+        client = GmailClient(email="[email protected]", secrets_json_string="{}")
+        with pytest.raises(GmailError):
+            client.get_raw_messages()
+
 
 class TestGetMessages:
     def test_it_returns_a_message_list(self, mocker, client, raw_incomplete_message):
@@ -50,6 +71,25 @@ def test_it_returns_a_raw_message(self, mocker, raw_complete_message):
         message = client.get_raw_message("123AAB")
         assert message == raw_complete_message
 
+    @pytest.mark.parametrize(
+        "error_code,exception_expected",
+        [(500, GmailError), (404, MessageNotFoundError)],
+    )
+    def test_it_encapsulates_gmail_exceptions(
+        self, mocker, error_code, exception_expected
+    ):
+        error_response = mocker.MagicMock(status=error_code)
+        mocker.patch(
+            "gmail_wrapper.client.GmailClient._make_client",
+            return_value=make_gmail_client(
+                mocker, get_effect=HttpError(error_response, b"Content")
+            ),
+        )
+
+        client = GmailClient(email="[email protected]", secrets_json_string="{}")
+        with pytest.raises(exception_expected):
+            client.get_raw_message("123AAB")
+
 
 class TestGetMessage:
     def test_it_returns_a_message(self, mocker, client, raw_complete_message):
@@ -77,6 +117,25 @@ def test_it_returns_a_raw_attachment_body(self, mocker, raw_attachment_body):
         )
         assert attachment_body == raw_attachment_body
 
+    @pytest.mark.parametrize(
+        "error_code,exception_expected",
+        [(500, GmailError), (404, AttachmentNotFoundError)],
+    )
+    def test_it_encapsulates_gmail_exceptions(
+        self, mocker, error_code, exception_expected
+    ):
+        error_response = mocker.MagicMock(status=error_code)
+        mocker.patch(
+            "gmail_wrapper.client.GmailClient._make_client",
+            return_value=make_gmail_client(
+                mocker, attachment_effect=HttpError(error_response, b"Content")
+            ),
+        )
+
+        client = GmailClient(email="[email protected]", secrets_json_string="{}")
+        with pytest.raises(exception_expected):
+            client.get_raw_attachment_body(id="CCX457", message_id="123AAB")
+
 
 class TestGetAttachmentBody:
     def test_it_returns_an_attachment_body(self, mocker, client, raw_attachment_body):
@@ -107,6 +166,27 @@ def test_it_modifies_and_return_a_raw_message(self, mocker, raw_complete_message
             body={"addLabelIds": ["processed"], "removeLabelIds": ["phishing"]},
         )
 
+    @pytest.mark.parametrize(
+        "error_code,exception_expected",
+        [(500, GmailError), (404, MessageNotFoundError)],
+    )
+    def test_it_encapsulates_gmail_exceptions(
+        self, mocker, error_code, exception_expected
+    ):
+        error_response = mocker.MagicMock(status=error_code)
+        mocker.patch(
+            "gmail_wrapper.client.GmailClient._make_client",
+            return_value=make_gmail_client(
+                mocker, modify_effect=HttpError(error_response, b"Content")
+            ),
+        )
+
+        client = GmailClient(email="[email protected]", secrets_json_string="{}")
+        with pytest.raises(exception_expected):
+            client.modify_raw_message(
+                id="CCX457", add_labels=["processed"], remove_labels=["phishing"]
+            )
+
 
 class TestModifyMessage:
     def test_it_returns_a_modified_message(self, client, mocker, raw_complete_message):

tests/utils.py

@@ -5,6 +5,10 @@ def make_gmail_client(
     attachment_return=None,
     modify_return=None,
     send_return=None,
+    list_effect=None,
+    get_effect=None,
+    attachment_effect=None,
+    modify_effect=None,
 ):
     return mocker.MagicMock(
         users=mocker.MagicMock(
@@ -13,23 +17,28 @@ def make_gmail_client(
                     return_value=mocker.MagicMock(
                         list=mocker.MagicMock(
                             return_value=mocker.MagicMock(
-                                execute=mocker.MagicMock(return_value=list_return)
+                                execute=mocker.MagicMock(
+                                    return_value=list_return, side_effect=list_effect
+                                )
                             )
                         ),
                         get=mocker.MagicMock(
                             return_value=mocker.MagicMock(
-                                execute=mocker.MagicMock(return_value=get_return)
+                                execute=mocker.MagicMock(
+                                    return_value=get_return, side_effect=get_effect
+                                )
                             )
                         ),
                         modify=mocker.MagicMock(
                             return_value=mocker.MagicMock(
-                                execute=mocker.MagicMock(return_value=modify_return)
+                                execute=mocker.MagicMock(
+                                    return_value=modify_return,
+                                    side_effect=modify_effect,
+                                )
                             )
                         ),
                         batchModify=mocker.MagicMock(
-                            return_value=mocker.MagicMock(
-                                execute=mocker.MagicMock()
-                            )
+                            return_value=mocker.MagicMock(execute=mocker.MagicMock())
                         ),
                         send=mocker.MagicMock(
                             return_value=mocker.MagicMock(
@@ -41,7 +50,8 @@ def make_gmail_client(
                                 get=mocker.MagicMock(
                                     return_value=mocker.MagicMock(
                                         execute=mocker.MagicMock(
-                                            return_value=attachment_return
+                                            return_value=attachment_return,
+                                            side_effect=attachment_effect,
                                         )
                                     )
                                 )