Fetch data from Gmail API with time range

gmail_fisher/api/gateway.py

@@ -65,11 +65,23 @@ def __authenticate(cls) -> Credentials:
 
 
 class GmailGateway:
-    """Maximum number of workers for thread pool executor"""
+    @staticmethod
+    def _get_search_query(sender_emails, keywords, start_time, end_time):
+        if start_time is None or end_time is None:
+            return f"from:{sender_emails} {keywords}"
+        else:
+            return (
+                f"from:{sender_emails} {keywords} after:{start_time} before:{end_time}"
+            )
 
     @classmethod
     def list_message_ids(
-        cls, sender_emails: str, keywords: str, max_results: int
+        cls,
+        sender_emails: str,
+        keywords: str,
+        max_results: int,
+        start_time: str = None,
+        end_time: str = None,
     ) -> Iterable[str]:
         """
         For a given sender email and comma-separated keywords, retrieve the matching
@@ -82,7 +94,7 @@ def list_message_ids(
             .messages()
             .list(
                 userId="me",
-                q=f"from:{sender_emails} {keywords}",
+                q=cls._get_search_query(sender_emails, keywords, start_time, end_time),
                 maxResults=max_results,
             )
             .execute(http=GmailClient.auth_http_request())
@@ -109,10 +121,12 @@ def get_email_messages(
         keywords: str,
         max_results: int,
         fetch_body: bool = False,
+        start_time: str = None,
+        end_time: str = None,
     ) -> Iterable[GmailMessage]:
         results = []
-        message_ids = GmailGateway.list_message_ids(
-            sender_emails, keywords, max_results
+        message_ids = cls.list_message_ids(
+            sender_emails, keywords, max_results, start_time, end_time
         )
 
         with ThreadPoolExecutor(max_workers=THREAD_POOL_MAX_WORKERS) as pool:
@@ -122,7 +136,7 @@ def get_email_messages(
             )
             with alive_bar(num_messages) as bar:
                 futures = [
-                    pool.submit(GmailGateway.get_message_detail, message_id, fetch_body)
+                    pool.submit(cls.get_message_detail, message_id, fetch_body)
                     for message_id in message_ids
                 ]
 
@@ -153,7 +167,7 @@ def get_message_detail(cls, message_id: str, fetch_body: bool) -> GmailMessage:
             .execute(http=GmailClient.auth_http_request())
         )
 
-        attachment_list = GmailGateway.get_message_attachments(
+        attachment_list = cls.get_message_attachments(
             get_message_result["payload"]
         )
         message_date = next(
@@ -173,7 +187,7 @@ def get_message_detail(cls, message_id: str, fetch_body: bool) -> GmailMessage:
         if not fetch_body:
             return message
 
-        message.body = GmailGateway.get_message_body(get_message_result["payload"])
+        message.body = cls.get_message_body(get_message_result["payload"])
 
         return message
 

gmail_fisher/parsers/food.py

@@ -22,9 +22,17 @@
 logger = get_logger(__name__)
 
 
-def apply_restaurant_filter(func):
+def apply_restaurant_filter(method):
+    """
+    Decorator method that wraps the 'get_restaurant' method from the food parsers
+    and applies all restaurant filters to the output of that method.
+
+    :param method: Pass the function that is being wrapped
+    :return: Filtered restaurant name
+    """
+
     def wrapper(*args, **kwargs):
-        return FoodExpenseParser.apply_restaurant_filters(func(*args, **kwargs))
+        return FoodExpenseParser.apply_restaurant_filters(method(*args, **kwargs))
 
     return wrapper
 
@@ -76,13 +84,17 @@ class BoltFoodParser(FoodExpenseParser):
     keywords: Final[str] = "Delivery from Bolt Food"
 
     @classmethod
-    def fetch_expenses(cls) -> Iterable[FoodExpense]:
+    def fetch_expenses(
+        cls, start_time: str = None, end_time: str = None
+    ) -> Iterable[FoodExpense]:
         print_header("🍕   Bolt Food")
         messages = GmailGateway.get_email_messages(
             sender_emails=cls.sender_email,
             keywords=cls.keywords,
             max_results=1000,
             fetch_body=True,
+            start_time=start_time,
+            end_time=end_time
         )
         return cls.parse_expenses_from_messages(messages)
 
@@ -164,7 +176,9 @@ class UberEatsParser(FoodExpenseParser):
     keywords: Final[str] = "Total"
 
     @classmethod
-    def fetch_expenses(cls) -> Iterable[FoodExpense]:
+    def fetch_expenses(
+        cls, start_time: str = None, end_time: str = None
+    ) -> Iterable[FoodExpense]:
         logger.info("Fetching UberEats food expenses")
         print_header("🍕   Uber Eats")
         messages = GmailGateway.get_email_messages(

gmail_fisher/services/food.py

@@ -14,22 +14,28 @@ def get_food_expenses() -> Iterable[FoodExpense]:
 
 
 def export_food_expenses(
-    service_type: FoodServiceType, output_path: Path, upload_s3: bool = False
+    service_type: FoodServiceType,
+    output_path: Path,
+    upload_s3: bool = False,
+    start_time: str = None,
+    end_time: str = None,
 ):
     logger.info(f"Exporting food expenses with {service_type=}, {output_path=}")
 
     if service_type is FoodServiceType.UBER_EATS:
         FoodExpenseParser.serialize_expenses_to_json_file(
-            expenses=UberEatsParser.fetch_expenses(), output_path=str(output_path)
+            expenses=UberEatsParser.fetch_expenses(start_time, end_time),
+            output_path=str(output_path),
         )
     elif service_type is FoodServiceType.BOLT_FOOD:
         FoodExpenseParser.serialize_expenses_to_json_file(
-            expenses=BoltFoodParser.fetch_expenses(), output_path=str(output_path)
+            expenses=BoltFoodParser.fetch_expenses(start_time, end_time),
+            output_path=str(output_path),
         )
     elif service_type is FoodServiceType.ALL:
         FoodExpenseParser.serialize_expenses_to_json_file(
-            expenses=list(BoltFoodParser.fetch_expenses())
-            + list(UberEatsParser.fetch_expenses()),
+            expenses=list(BoltFoodParser.fetch_expenses(start_time, end_time))
+            + list(UberEatsParser.fetch_expenses(start_time, end_time)),
             output_path=str(output_path),
         )
     else:

gmail_fisher/utils/config.py

@@ -15,4 +15,5 @@
 AUTH_PATH: Final[Path] = Path("auth/")
 
 # CONCURRENCY
+# Maximum number of workers for thread pool executor
 THREAD_POOL_MAX_WORKERS: Final[int] = 200

pyproject.toml

@@ -32,6 +32,7 @@ boto3 = "^1.21.13"
 python-dotenv = "^0.20.0"
 pdfplumber = "^0.6.0"
 alive-progress = "^3.0.1"
+pytest-mock = "^3.10.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.0.1"

tests/conftest.py

@@ -5,6 +5,24 @@
 from gmail_fisher.data.models import GmailMessage
 
 
+@pytest.fixture
+def gmail_api_stub(mocker):
+    # Set up the desired response from the API
+    response = {
+        "messages": [
+            {"id": "message_id_1", "threadId": "thread_id_1"},
+            {"id": "message_id_2", "threadId": "thread_id_2"},
+        ]
+    }
+
+    # Stub the requests to the Gmail API
+    mocker.patch(
+        "googleapiclient.discovery.build"
+    ).return_value.users.return_value.messages.return_value.list.return_value.execute.return_value = (
+        response
+    )
+
+
 @pytest.fixture
 def bolt_email_html_body() -> str:
     return """<!DOCTYPE HTML>

tests/test_gmail_gateway.py

@@ -0,0 +1,8 @@
+from gmail_fisher.api.gateway import GmailGateway
+
+
+def test_get_email_messages(gmail_api_stub):
+    messages = GmailGateway.list_message_ids(
+        sender_emails="portugal-food@bolt.eu", keywords="Delivery", max_results=10
+    )
+    print(len(messages))