Merge pull request #3 from loadsmart/labelSupport

Support labels modification
loadsmart · Aug 28, 2019 · ffabf53 · ffabf53
2 parents 3a99514 + 73e0b16
commit ffabf53
Show file tree

Hide file tree

Showing 7 changed files with 105 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -1,10 +1,9 @@
 # Gmail Wrapper
 
-[![CircleCI](https://circleci.com/gh/loadsmart/gmail-wrapper/tree/master.svg?style=svg)](https://circleci.com/gh/loadsmart/gmail-wrapper/tree/master)
-
+[![CircleCI](https://circleci.com/gh/loadsmart/gmail-wrapper/tree/master.svg?style=svg&circle-token=110f54407b50c79865fe1f9b4352e213bc68504b)](https://circleci.com/gh/loadsmart/gmail-wrapper/tree/master)
 [![codecov](https://codecov.io/gh/loadsmart/gmail-wrapper/branch/master/graph/badge.svg?token=Ciq3QScb0L)](https://codecov.io/gh/loadsmart/gmail-wrapper)
 
-Because scrapping Gmail data doesn't have to be a pain.
+Because scrapping Gmail data doesn't have to be a [pain](https://googleapis.github.io/google-api-python-client/docs/dyn/gmail_v1.html).
 
 ## Installing
 
@@ -53,25 +52,26 @@ client = GmailClient(email_account, secrets_json_string=credentials_string, scop
 ```python
 import sys
 
-def print_message_data(message):
+query = "filename:pdf label:friends" # Check Gmail query docs: https://support.google.com/mail/answer/7190
+messages = client.get_messages(query=query, limit=10)
+for message in messages:
     print("-- MESSAGE {} --".format(message.id))
     print("SUBJECT: {}".format(message.subject))
     print("DATE: {}".format(message.date))
     for attachment in message.attachments:
         print("\t-- ATTACHMENT {} --".format(attachment.id))
         print("\t\tFILENAME: {}".format(attachment.filename))
         print("\t\tDECODED SIZE: {}".format(sys.getsizeof(attachment.content)))
+```
 
-# Get a message list
-query = "filename:pdf label:friends" # Check Gmail query docs: https://support.google.com/mail/answer/7190
-messages = client.get_messages(query=query, limit=10)
-for message in messages:
-    print_message_data(message)
+- Modify message labels
 
-# ...Or get a single message
+```python
 message_id = "..."
 message = client.get_message(message_id)
-print_message_data(message)
+print(message.labels) # ["foo", "bar"]
+message.modify(add_labels=["processed"], remove_labels=["foo"]) # Beware that you'll need proper scopes
+print(message.labels) # ["bar", "processed"]
 ```
 
 ## Need more help?

diff --git a/gmail_wrapper/client.py b/gmail_wrapper/client.py
@@ -59,6 +59,21 @@ def get_message(self, 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 [],
+            },
+        ).execute()
+
+    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 get_raw_attachment_body(self, id, message_id):
         return (
             self._messages_resource()

diff --git a/gmail_wrapper/entities.py b/gmail_wrapper/entities.py
@@ -66,6 +66,13 @@ def headers(self):
             headers[header["name"]] = header["value"]
         return headers
 
+    @property
+    def labels(self):
+        if "labelIds" not in self._raw:
+            self._raw = self._client.get_raw_message(self.id)
+
+        return self._raw["labelIds"]
+
     @property
     def id(self):
         return self._raw.get("id")
@@ -88,5 +95,10 @@ def attachments(self):
             if part["filename"]
         ]
 
+    def modify(self, add_labels=None, remove_labels=None):
+        self._raw = self._client.modify_raw_message(
+            self.id, add_labels=add_labels, remove_labels=remove_labels
+        )
+
     def __str__(self):
         return "Gmail message: {}".format(self.id)
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -26,7 +26,7 @@ def secrets_string():
 
 @pytest.fixture
 def raw_incomplete_message():
-    return {"id": 123, "internalDate": "1566398665"}
+    return {"id": "123AAB", "internalDate": "1566398665"}
 
 
 @pytest.fixture

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -86,3 +86,35 @@ def test_it_returns_an_attachment_body(self, mocker, client, raw_attachment_body
         mocked_get_raw_attachment_body.assert_called_once_with("CCX457", "123AAB")
         assert isinstance(attachment_body, AttachmentBody)
         assert attachment_body.id == raw_attachment_body["attachmentId"]
+
+
+class TestModifyRawMessage:
+    def test_it_modifies_and_return_a_raw_message(self, mocker, raw_complete_message):
+        mocker.patch(
+            "gmail_wrapper.client.GmailClient._make_client",
+            return_value=make_gmail_client(mocker, modify_return=raw_complete_message),
+        )
+        client = GmailClient(email="[email protected]", secrets_json_string="{}")
+        modified_message = client.modify_raw_message(
+            id="CCX457", add_labels=["processed"], remove_labels=["phishing"]
+        )
+        assert modified_message == raw_complete_message
+        client._messages_resource().modify.assert_called_once_with(
+            userId="[email protected]",
+            id="CCX457",
+            body={"addLabelIds": ["processed"], "removeLabelIds": ["phishing"]},
+        )
+
+
+class TestModifyMessage:
+    def test_it_returns_a_modified_message(self, client, mocker, raw_complete_message):
+        mocked_modify_raw_message = mocker.patch(
+            "gmail_wrapper.client.GmailClient.modify_raw_message",
+            return_value=raw_complete_message,
+        )
+        modified_message = client.modify_message(
+            id="CCX457", add_labels=["foo"], remove_labels=["bar"]
+        )
+        mocked_modify_raw_message.assert_called_once_with("CCX457", ["foo"], ["bar"])
+        assert isinstance(modified_message, Message)
+        assert modified_message.id == raw_complete_message["id"]
diff --git a/tests/test_entities.py b/tests/test_entities.py
@@ -35,6 +35,17 @@ def test_it_fetch_additional_information_when_needed(
         )
         mocked_get_raw_message.assert_called_once_with(raw_incomplete_message["id"])
 
+    def test_get_labels(
+        self, mocker, raw_complete_message, client, raw_incomplete_message
+    ):
+        mocked_get_raw_message = mocker.patch(
+            "gmail_wrapper.client.GmailClient.get_raw_message",
+            return_value=raw_complete_message,
+        )
+        incomplete_message = Message(client, raw_incomplete_message)
+        assert incomplete_message.labels == raw_complete_message["labelIds"]
+        mocked_get_raw_message.assert_called_once_with(raw_incomplete_message["id"])
+
     def test_as_str(self, client, raw_incomplete_message):
         message = Message(client, raw_incomplete_message)
         assert str(message) == "Gmail message: {}".format(raw_incomplete_message["id"])
@@ -48,6 +59,22 @@ def test_it_generates_attachment_objects(self, client, raw_complete_message):
             ]
         )
 
+    def test_it_modifies_a_message(
+        self, mocker, client, raw_incomplete_message, raw_complete_message
+    ):
+        mocked_modify_raw_message = mocker.patch(
+            "gmail_wrapper.client.GmailClient.modify_raw_message",
+            return_value=raw_complete_message,
+        )
+        message = Message(client, raw_incomplete_message)
+        message.modify(add_labels=["phishing"], remove_labels=["oklahoma"])
+        mocked_modify_raw_message.assert_called_once_with(
+            raw_incomplete_message["id"],
+            add_labels=["phishing"],
+            remove_labels=["oklahoma"],
+        )
+        assert message.labels == raw_complete_message["labelIds"]
+
 
 class TestAttachment:
     def test_it_has_basic_properties_without_additional_fetch(

diff --git a/tests/utils.py b/tests/utils.py
@@ -1,5 +1,5 @@
 def make_gmail_client(
-    mocker, list_return=None, get_return=None, attachment_return=None
+        mocker, list_return=None, get_return=None, attachment_return=None, modify_return=None
 ):
     return mocker.MagicMock(
         users=mocker.MagicMock(
@@ -16,6 +16,11 @@ def make_gmail_client(
                                 execute=mocker.MagicMock(return_value=get_return)
                             )
                         ),
+                        modify=mocker.MagicMock(
+                            return_value=mocker.MagicMock(
+                                execute=mocker.MagicMock(return_value=modify_return)
+                            )
+                        ),
                         attachments=mocker.MagicMock(
                             return_value=mocker.MagicMock(
                                 get=mocker.MagicMock(
@@ -31,4 +36,4 @@ def make_gmail_client(
                 )
             )
         )
-    )
+    )