diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 1d744a9..56d5b7c 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -39,3 +39,40 @@ jobs: run: | pip install --no-deps . coverage run --omit gmailsorter/_version.py -m unittest discover tests + + imap-integration: + runs-on: ubuntu-latest + + services: + greenmail: + image: greenmail/standalone:2.1.11 + env: + GREENMAIL_OPTS: >- + -Dgreenmail.setup.test.smtp + -Dgreenmail.setup.test.imap + -Dgreenmail.hostname=0.0.0.0 + -Dgreenmail.users=testuser:secret@example.test + ports: + - 3025:3025 + - 3143:3143 + + env: + TEST_SMTP_HOST: localhost + TEST_SMTP_PORT: "3025" + TEST_IMAP_HOST: localhost + TEST_IMAP_PORT: "3143" + TEST_IMAP_USERNAME: testuser + TEST_EMAIL: testuser@example.test + TEST_EMAIL_PASSWORD: secret + # fail this job rather than skip it if the greenmail service never came up + IMAP_INTEGRATION_REQUIRED: "true" + + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - name: Install package + run: pip install . + - name: Run IMAP integration test + run: python -m unittest tests.test_imap_service_integration -v diff --git a/docs/source/architecture.md b/docs/source/architecture.md index d797a14..93214ca 100644 --- a/docs/source/architecture.md +++ b/docs/source/architecture.md @@ -9,9 +9,12 @@ machine learning knowledge required. If you are looking for setup instructions i Regardless of whether you use the hosted [gmailsorter.com](https://gmailsorter.com) service, the Docker container or the plain Python package, `gmailsorter` is built from the same three building blocks: -* **Your Google Mail account** - the source of truth for your emails and labels, accessed exclusively through the - official [Gmail API](https://developers.google.com/gmail/api/guides). `gmailsorter` never reads your mailbox - through any other channel and never stores your Google password. +* **Your email account** - the source of truth for your emails and labels, accessed either through the official + [Gmail API](https://developers.google.com/gmail/api/guides) or, for any other IMAP-capable provider, through a + plain IMAP connection. `gmailsorter` never stores your Google password, and for IMAP accounts the password you + provide is used only to log in - it is not persisted anywhere. When talking to a plain IMAP server, each mailbox + folder plays the role a Gmail label plays throughout the rest of this page - "moving" an email between labels + means moving it between IMAP folders. * **A local database** - a SQL database (SQLite by default, though any database supported by [SQLAlchemy](https://www.sqlalchemy.org/) works) that keeps a private copy of your email metadata, your login token and your trained models. In the Docker container and the plain Python package this database lives entirely diff --git a/docs/source/developer.md b/docs/source/developer.md index 7a35901..77b6b19 100644 --- a/docs/source/developer.md +++ b/docs/source/developer.md @@ -60,6 +60,34 @@ a selected label `"MyLabel"`. Then reloads the machine learning model from the l the correct labels for these emails. The `recommendation_ratio` defines the level of certainty required to actually move the email, with `0.9` equalling a certainty of 90%. +## IMAP accounts +`gmailsorter` also supports plain IMAP accounts (username and password, e.g. an app +password), for mail servers other than Google Mail. Import the `Imap` class instead of +`Gmail`: +``` +from gmailsorter import Imap +``` +``` +imap = Imap( + host="imap.example.com", + port=993, + username="user@example.com", + password="app-password", + connection_str="sqlite:////absolute/path/to/email.db", +) +``` +`Imap` exposes the exact same `update_database()`, `get_all_emails_in_database()` and +`filter_messages_from_server()` methods as `Gmail` - the only difference is that IMAP +folders play the role Gmail labels play elsewhere in this document: each folder is +treated as one label, and moving an email means moving it from one IMAP folder to +another. A command line interface is also available as `gmailsorter-imap`, reading the +account password from an environment variable (`IMAP_PASSWORD` by default) rather than +accepting it as a command line argument: +``` +export IMAP_PASSWORD=app-password +gmailsorter-imap --host imap.example.com --username user@example.com -d sqlite:///email.db -u +``` + ## Future directions The current machine learning model is limited in the precision and memory usage. So there is a great interest to replace it with a computationally more efficient model. All suggestions and feedback are welcome. Beyond the optimization of the diff --git a/gmailsorter/__init__.py b/gmailsorter/__init__.py index 49301a6..a9f2fba 100644 --- a/gmailsorter/__init__.py +++ b/gmailsorter/__init__.py @@ -1,6 +1,6 @@ -from gmailsorter.local import Gmail, load_client_secrets_file +from gmailsorter.local import Gmail, Imap, load_client_secrets_file from . import _version __version__: str = _version.__version__ -__all__ = ["Gmail", "load_client_secrets_file"] +__all__ = ["Gmail", "Imap", "load_client_secrets_file"] diff --git a/gmailsorter/__main__.py b/gmailsorter/__main__.py index 36615a8..f68dfe6 100644 --- a/gmailsorter/__main__.py +++ b/gmailsorter/__main__.py @@ -79,7 +79,9 @@ def command_line_parser() -> None: ) elif args.label: gmail.filter_messages_from_server( - label=args.label, recommendation_ratio=0.9 + label=args.label, + recommendation_ratio=0.9, + label_prefix="labels_Label_", ) else: parser.print_help() diff --git a/gmailsorter/base/mail.py b/gmailsorter/base/mail.py new file mode 100644 index 0000000..9b18f20 --- /dev/null +++ b/gmailsorter/base/mail.py @@ -0,0 +1,329 @@ +from abc import ABC, abstractmethod + +import pandas +from tqdm import tqdm + +from gmailsorter.ml import ( + encode_df_for_machine_learning, + fit_machine_learning_models, + get_predictions_from_machine_learning_models, +) + + +class AbstractMailBox(ABC): + def __init__( + self, + mail_service, + database_email=None, + database_ml=None, + user_id="me", + db_user_id=1, + email_download_format="metadata", + ): + """ + Shared fetch-store-train-predict-move loop for a mailbox backend, independent of + whether the backend is the Gmail API or a plain IMAP connection. + + Args: + mail_service: backend-specific connection object (Gmail API service resource, + imaplib connection, ...) + database_email (gmailsorter.base.database.DatabaseInterface): SQLalchemy interface for email database + database_ml (gmailsorter.ml.database.DatabaseInterface): SQLalchemy interface for machine learning database + user_id (str): backend-specific user identifier + db_user_id (int): Default 1 - set a user id when sharing a database with multiple users + email_download_format (str): backend-specific download format hint + """ + self._service = mail_service + self._db_email = database_email + self._db_ml = database_ml + self._db_user_id = db_user_id + self._userid = user_id + self._email_download_format = email_download_format + self._label_dict = self._get_label_translate_dict() + self._label_dict_inverse = {v: k for k, v in self._label_dict.items()} + + @property + def labels(self): + return list(self._label_dict.keys()) + + def download_emails_for_label(self, label): + """ + Download emails for a specific label + + Args: + label (str): label to download emails for + + Returns: + pandas.DataFrame: Email content for the downloaded emails + """ + return self._download_messages_to_dataframe( + message_id_lst=self._search_email_on_server( + label_lst=[label], only_message_ids=True + ) + ) + + def filter_messages_from_server( + self, + label, + recommendation_ratio=0.9, + label_prefix: str = "labels_", + ): + """ + Filter new emails based on machine learning model recommendations. + + Args: + label (str): Email label to filter for + recommendation_ratio (float): Only accept recommendation above this ratio (0 0: + model_reload_dict, feature_reload_lst = self._db_ml.load_models() + df_partial_features = encode_df_for_machine_learning( + df=df_partial, + feature_lst=feature_reload_lst, + label_lst=list(model_reload_dict.keys()), + return_labels=False, + label_prefix=label_prefix, + ) + df_partial_features = df_partial_features.reindex( + sorted(df_partial_features.columns), axis=1 + ) + model_recommendation_dict = get_predictions_from_machine_learning_models( + df_features=df_partial_features, + model_dict=model_reload_dict, + recommendation_ratio=recommendation_ratio, + ) + self._move_emails( + move_email_dict=model_recommendation_dict, label_to_ignore=label + ) + + def fit_machine_learning_model_to_database( + self, + n_estimators=100, + max_features=400, + random_state=42, + bootstrap=True, + include_deleted=False, + max_workers=None, + ): + """ + Fit machine learning models to emails stored in database and afterwards store machine learning models in + database. + + Args: + n_estimators (int): Number of estimators + max_features (int): Number of features + random_state (int): Random state + bootstrap (boolean): Whether bootstrap samples are used when building trees. If False, the whole dataset is + used to build each tree. (default: true) + include_deleted (bool): Flag to include deleted emails - default False + max_workers (int): maximum number of workers for the machine learning models + """ + df_all = self.get_all_emails_in_database(include_deleted=include_deleted) + df_all_features, df_all_labels = encode_df_for_machine_learning( + df=df_all, feature_lst=[], label_lst=[], return_labels=True + ) + df_all_features = df_all_features.loc[ + :, ~df_all_features.columns.duplicated() + ].copy() + df_all_features = df_all_features.reindex( + sorted(df_all_features.columns), axis=1 + ) + model_dict = fit_machine_learning_models( + df_all_features=df_all_features, + df_all_labels=df_all_labels, + n_estimators=n_estimators, + max_features=max_features, + random_state=random_state, + bootstrap=bootstrap, + max_workers=max_workers, + ) + self._db_ml.store_models( + model_dict=model_dict, + feature_lst=df_all_features.columns.values.tolist(), + user_id=self._db_user_id, + commit=True, + ) + + def get_all_emails_in_database(self, include_deleted=False): + """ + Get all emails stored in the local database + + Args: + include_deleted (bool): Flag to include deleted emails - default False + + Returns: + pandas.DataFrame: With all emails and the corresponding information + """ + return self._db_email.get_all_emails( + include_deleted=include_deleted, user_id=self._db_user_id + ) + + def update_database(self, quick=False, label_lst=None, email_format=None): + """ + Update local email database + + Args: + quick (boolean): Only add new emails, do not update existing labels - by default: False + label_lst (list): list of labels to be searched + email_format (str/None): Email format to download + """ + if label_lst is None: + label_lst = [] + if self._db_email is not None: + message_id_lst = self._search_email_on_server( + label_lst=label_lst, only_message_ids=True + ) + ( + new_messages_lst, + message_label_updates_lst, + deleted_messages_lst, + ) = self._db_email.get_labels_to_update( + message_id_lst=message_id_lst, user_id=self._db_user_id + ) + if not quick: + self._db_email.mark_emails_as_deleted( + message_id_lst=deleted_messages_lst, user_id=self._db_user_id + ) + self._db_email.update_labels( + message_id_lst=message_label_updates_lst, + message_meta_lst=self._get_labels_for_emails( + message_id_lst=message_label_updates_lst + ), + user_id=self._db_user_id, + ) + self._store_emails_in_database( + message_id_lst=new_messages_lst, email_format=email_format + ) + + def _download_messages_to_dataframe(self, message_id_lst, email_format=None): + """ + Download a list of messages based on their email IDs and store the content in a pandas.DataFrame. + + Args: + message_id_lst (list): list of emails IDs + email_format (str): Email format to download - default: "full" + + Returns: + pandas.DataFrame: pandas.DataFrame which contains the rendered emails + """ + return pandas.DataFrame( + [ + message + for message in [ + self._parse_message( + message=self._get_message_detail( + message_id=message_id, + email_format=email_format, + metadata_headers=[], + ) + ) + for message_id in tqdm( + iterable=message_id_lst, desc="Download messages to DataFrame" + ) + ] + if message is not None + ] + ) + + def _get_labels_for_emails(self, message_id_lst): + """ + Get labels for a list of emails + + Args: + message_id_lst (list): list of emails IDs + + Returns: + list: Nested list of email labels for each email + """ + return [ + self._get_labels_for_email(message_id=message_id) + for message_id in tqdm( + iterable=message_id_lst, desc="Get labels for emails" + ) + ] + + def _move_emails(self, move_email_dict, label_to_ignore): + label_existing = self._label_dict[label_to_ignore] + for message_id, label_add in tqdm( + iterable=move_email_dict.items(), desc="Move emails" + ): + if label_add is not None and label_add != label_existing: + self._modify_message_labels( + message_id=message_id, + label_id_remove_lst=[label_existing], + label_id_add_lst=[label_add], + ) + + def _store_emails_in_database(self, message_id_lst, email_format=None): + df = self._download_messages_to_dataframe( + message_id_lst=message_id_lst, email_format=email_format + ) + if len(df) > 0: + self._db_email.store_dataframe(df=df, user_id=self._db_user_id) + + @abstractmethod + def _search_email_on_server( + self, query_string="", label_lst=None, only_message_ids=False + ): + """ + Search emails either by a specific query or optionally limit your search to a list of labels + + Args: + query_string (str): query string to search for + label_lst (list): list of labels to be searched + only_message_ids (bool): return only the email IDs not the thread IDs - default: false + + Returns: + list: list of message ids (or backend-specific list items) matching the search + """ + + @abstractmethod + def _get_message_detail(self, message_id, email_format=None, metadata_headers=None): + """ + Get the raw, backend-specific representation of a single email message. + + Args: + message_id (str): id used by this backend to uniquely identify the email + email_format (str/None): backend-specific format hint + metadata_headers (list): backend-specific list of metadata headers + + Returns: + The backend-specific raw message representation, passed on to `_parse_message`. + """ + + @abstractmethod + def _get_label_translate_dict(self): + """ + Returns: + dict: mapping of label/folder display name to the backend-specific label/folder id + """ + + @abstractmethod + def _modify_message_labels( + self, message_id, label_id_remove_lst=None, label_id_add_lst=None + ): + """ + Apply a label/folder change to a single email message. + """ + + @abstractmethod + def _get_labels_for_email(self, message_id): + """ + Args: + message_id (str): id used by this backend to uniquely identify the email + + Returns: + list: list of labels/folders currently assigned to the email + """ + + @abstractmethod + def _parse_message(self, message): + """ + Args: + message: the backend-specific raw message representation returned by `_get_message_detail` + + Returns: + dict/None: the common gmailsorter email dict (see `gmailsorter.base.message.AbstractMessage.to_dict`), + or None if the message could not be parsed + """ diff --git a/gmailsorter/base/message.py b/gmailsorter/base/message.py index a58ca48..90b2843 100644 --- a/gmailsorter/base/message.py +++ b/gmailsorter/base/message.py @@ -1,11 +1,35 @@ from abc import ABC, abstractmethod from datetime import datetime +from html.parser import HTMLParser +from io import StringIO from typing import Any _MAX_DATE_COMMAS = 2 _DATE_HYPHEN_COUNT = 2 +# https://stackoverflow.com/questions/753052/strip-html-from-strings-in-python +class _MLStripper(HTMLParser): + def __init__(self): + super().__init__() + self.reset() + self.strict = False + self.convert_charrefs = True + self.text = StringIO() + + def handle_data(self, d): + self.text.write(d) + + def get_data(self): + return self.text.getvalue() + + +def strip_html_tags(html): + stripper = _MLStripper() + stripper.feed(html) + return stripper.get_data() + + def email_date_converter(email_date: Any) -> datetime | None: if not isinstance(email_date, str): return None diff --git a/gmailsorter/daemon/daemon.py b/gmailsorter/daemon/daemon.py index 006e11b..ad67d09 100644 --- a/gmailsorter/daemon/daemon.py +++ b/gmailsorter/daemon/daemon.py @@ -134,6 +134,7 @@ def iterate_over_users( gmail.filter_messages_from_server( label=MAILSORT_LABEL, recommendation_ratio=recommendation_ratio, + label_prefix="labels_Label_", ) except HttpError: update_task_status( diff --git a/gmailsorter/google/mail.py b/gmailsorter/google/mail.py index 9018508..997f8ca 100644 --- a/gmailsorter/google/mail.py +++ b/gmailsorter/google/mail.py @@ -1,21 +1,17 @@ from typing import Any -import pandas from googleapiclient.discovery import Resource from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker -from tqdm import tqdm from gmailsorter.base import get_email_database from gmailsorter.base.database import DatabaseInterface as EmailDatabaseInterface +from gmailsorter.base.mail import AbstractMailBox from gmailsorter.google.database import DatabaseInterface as TokenDatabaseInterface from gmailsorter.google.database import get_token_database from gmailsorter.google.message import get_email_dict from gmailsorter.ml import ( - encode_df_for_machine_learning, - fit_machine_learning_models, get_machine_learning_database, - get_predictions_from_machine_learning_models, ) from gmailsorter.ml.database import MachineLearningDatabase @@ -24,7 +20,7 @@ ] -class GoogleMailBase: +class GoogleMailBase(AbstractMailBox): def __init__( self, google_mail_service: Resource, @@ -47,248 +43,15 @@ def __init__( db_user_id (int): Default 1 - set a user id when sharing a database with multiple users email_download_format (str): API response format [full, metadata] """ - self._service = google_mail_service - self._db_email = database_email - self._db_ml = database_ml self._db_token = database_token - self._db_user_id = db_user_id - self._userid = user_id - self._email_download_format = email_download_format - self._label_dict = self._get_label_translate_dict() - self._label_dict_inverse = {v: k for k, v in self._label_dict.items()} - - @property - def labels(self) -> list[str]: - return list(self._label_dict.keys()) - - def download_emails_for_label(self, label: str) -> pandas.DataFrame: - """ - Download emails for a specific label - - Args: - label (str): label to download emails for - - Returns: - pandas.DataFrame: Email content for the downloaded emails - """ - return self._download_messages_to_dataframe( - message_id_lst=self._search_email_on_server( - label_lst=[label], only_message_ids=True - ) - ) - - def filter_messages_from_server( - self, - label: str, - recommendation_ratio: float = 0.9, - ) -> None: - """ - Filter new emails based on machine learning model recommendations. - - Args: - label (str): Email label to filter for - recommendation_ratio (float): Only accept recommendation above this ratio (0 0: - model_reload_dict, feature_reload_lst = self._db_ml.load_models() - df_partial_features = encode_df_for_machine_learning( - df=df_partial, - feature_lst=feature_reload_lst, - label_lst=list(model_reload_dict.keys()), - return_labels=False, - label_prefix="labels_Label_", - ) - df_partial_features = df_partial_features.reindex( - sorted(df_partial_features.columns), axis=1 - ) - model_recommendation_dict = get_predictions_from_machine_learning_models( - df_features=df_partial_features, - model_dict=model_reload_dict, - recommendation_ratio=recommendation_ratio, - ) - self._move_emails( - move_email_dict=model_recommendation_dict, label_to_ignore=label - ) - - def fit_machine_learning_model_to_database( - self, - n_estimators: int = 100, - max_features: int = 400, - random_state: int = 42, - bootstrap: bool = True, - include_deleted: bool = False, - max_workers: int | None = None, - ): - """ - Fit machine learning models to emails stored in database and afterwards store machine learning models in - database. - - Args: - n_estimators (int): Number of estimators - max_features (int): Number of features - random_state (int): Random state - bootstrap (boolean): Whether bootstrap samples are used when building trees. If False, the whole dataset is - used to build each tree. (default: true) - include_deleted (bool): Flag to include deleted emails - default False - max_workers (int): maximum number of workers for the machine learning models - """ - df_all = self.get_all_emails_in_database(include_deleted=include_deleted) - df_all_features, df_all_labels = encode_df_for_machine_learning( - df=df_all, - feature_lst=[], - label_lst=[], - return_labels=True, - label_prefix="labels_Label_", - ) - df_all_features = df_all_features.loc[ - :, ~df_all_features.columns.duplicated() - ].copy() - df_all_features = df_all_features.reindex( - sorted(df_all_features.columns), axis=1 - ) - model_dict = fit_machine_learning_models( - df_all_features=df_all_features, - df_all_labels=df_all_labels, - n_estimators=n_estimators, - max_features=max_features, - random_state=random_state, - bootstrap=bootstrap, - max_workers=max_workers, + super().__init__( + mail_service=google_mail_service, + database_email=database_email, + database_ml=database_ml, + user_id=user_id, + db_user_id=db_user_id, + email_download_format=email_download_format, ) - self._db_ml.store_models( - model_dict=model_dict, - feature_lst=df_all_features.columns.values.tolist(), - user_id=self._db_user_id, - commit=True, - ) - - def get_all_emails_in_database( - self, include_deleted: bool = False - ) -> pandas.DataFrame: - """ - Get all emails stored in the local database - - Args: - include_deleted (bool): Flag to include deleted emails - default False - - Returns: - pandas.DataFrame: With all emails and the corresponding information - """ - return self._db_email.get_all_emails( - include_deleted=include_deleted, user_id=self._db_user_id - ) - - def update_database( - self, - quick: bool = False, - label_lst: list[str] | None = None, - email_format: str | None = None, - ) -> None: - """ - Update local email database - - Args: - quick (boolean): Only add new emails, do not update existing labels - by default: False - label_lst (list): list of labels to be searched - email_format (str/None): Email format to download - """ - if label_lst is None: - label_lst = [] - if self._db_email is not None: - message_id_lst = self._search_email_on_server( - label_lst=label_lst, only_message_ids=True - ) - ( - new_messages_lst, - message_label_updates_lst, - deleted_messages_lst, - ) = self._db_email.get_labels_to_update( - message_id_lst=message_id_lst, user_id=self._db_user_id - ) - if not quick: - self._db_email.mark_emails_as_deleted( - message_id_lst=deleted_messages_lst, user_id=self._db_user_id - ) - self._db_email.update_labels( - message_id_lst=message_label_updates_lst, - message_meta_lst=self._get_labels_for_emails( - message_id_lst=message_label_updates_lst - ), - user_id=self._db_user_id, - ) - self._store_emails_in_database( - message_id_lst=new_messages_lst, email_format=email_format - ) - - def _download_messages_to_dataframe( - self, message_id_lst: list[str], email_format: str | None = None - ) -> pandas.DataFrame: - """ - Download a list of messages based on their email IDs and store the content in a pandas.DataFrame. - - Args: - message_id_lst (list): list of emails IDs - email_format (str): Email format to download - default: "full" - - Returns: - pandas.DataFrame: pandas.DataFrame which contains the rendered emails - """ - return pandas.DataFrame( - [ - message - for message in [ - get_email_dict( - message=self._get_message_detail( - message_id=message_id, - email_format=email_format, - metadata_headers=[], - ) - ) - for message_id in tqdm( - iterable=message_id_lst, desc="Download messages to DataFrame" - ) - ] - if message is not None - ] - ) - - def _get_labels_for_email(self, message_id: str) -> list[str]: - """ - Get labels for email - - Args: - message_id (str): email ID - - Returns: - list: List of email labels - """ - message_dict = self._get_message_detail( - message_id=message_id, - email_format="metadata", - metadata_headers=["labelIds"], - ) - if "labelIds" in message_dict: - return message_dict["labelIds"] - else: - return [] - - def _get_labels_for_emails(self, message_id_lst: list[str]) -> list[list[str]]: - """ - Get labels for a list of emails - - Args: - message_id_lst (list): list of emails IDs - - Returns: - list: Nested list of email labels for each email - """ - return [ - self._get_labels_for_email(message_id=message_id) - for message_id in tqdm( - iterable=message_id_lst, desc="Get labels for emails" - ) - ] def _get_label_translate_dict(self) -> dict[str, str]: results = self._service.users().labels().list(userId=self._userid).execute() @@ -390,20 +153,6 @@ def _modify_message_labels( userId=self._userid, id=message_id, body=body_dict ).execute() - def _move_emails( - self, move_email_dict: dict[str, str | None], label_to_ignore: str - ) -> None: - label_existing = self._label_dict[label_to_ignore] - for message_id, label_add in tqdm( - iterable=move_email_dict.items(), desc="Move emails" - ): - if label_add is not None and label_add != label_existing: - self._modify_message_labels( - message_id=message_id, - label_id_remove_lst=[label_existing], - label_id_add_lst=[label_add], - ) - def _search_email_on_server( self, query_string: str = "", @@ -432,14 +181,28 @@ def _search_email_on_server( else: return [d["id"] for d in message_id_lst] - def _store_emails_in_database( - self, message_id_lst: list[str], email_format: str | None = None - ) -> None: - df = self._download_messages_to_dataframe( - message_id_lst=message_id_lst, email_format=email_format + def _get_labels_for_email(self, message_id: str) -> list[str]: + """ + Get labels for email + + Args: + message_id (str): email ID + + Returns: + list: List of email labels + """ + message_dict = self._get_message_detail( + message_id=message_id, + email_format="metadata", + metadata_headers=["labelIds"], ) - if len(df) > 0: - self._db_email.store_dataframe(df=df, user_id=self._db_user_id) + if "labelIds" in message_dict: + return message_dict["labelIds"] + else: + return [] + + def _parse_message(self, message): + return get_email_dict(message=message) @staticmethod def _create_databases(connection_str: str) -> _DatabaseTriple: diff --git a/gmailsorter/google/message.py b/gmailsorter/google/message.py index 56bb68d..a259827 100644 --- a/gmailsorter/google/message.py +++ b/gmailsorter/google/message.py @@ -1,26 +1,12 @@ import base64 from datetime import datetime -from html.parser import HTMLParser -from io import StringIO from typing import Any -from gmailsorter.base.message import AbstractMessage, email_date_converter - - -# https://stackoverflow.com/questions/753052/strip-html-from-strings-in-python -class MLStripper(HTMLParser): - def __init__(self) -> None: - super().__init__() - self.reset() - self.strict = False - self.convert_charrefs = True - self.text = StringIO() - - def handle_data(self, d: str) -> None: - self.text.write(d) - - def get_data(self) -> str: - return self.text.getvalue() +from gmailsorter.base.message import ( + AbstractMessage, + email_date_converter, + strip_html_tags, +) def get_email_dict(message: dict[str, Any]) -> dict[str, Any] | None: @@ -102,7 +88,7 @@ def _get_parts_content(self, message_parts: list[dict[str, Any]]) -> str | None: message_parts=message_parts[content_types.index("text/plain")] ) elif "text/html" in content_types: - return self._strip_tags( + return strip_html_tags( html=self._get_email_body( message_parts=message_parts[content_types.index("text/html")] ) @@ -140,12 +126,6 @@ def _get_email_body(message_parts: dict[str, Any]) -> str: else: return "" - @staticmethod - def _strip_tags(html: str) -> str: - s = MLStripper() - s.feed(html) - return s.get_data() - @staticmethod def _get_email_address(email: str) -> str: email_split = email.split("<") diff --git a/gmailsorter/imap/__init__.py b/gmailsorter/imap/__init__.py new file mode 100644 index 0000000..4bb407f --- /dev/null +++ b/gmailsorter/imap/__init__.py @@ -0,0 +1,4 @@ +from gmailsorter.imap.authentication import create_service +from gmailsorter.imap.mail import ImapMailBase + +__all__ = ["create_service", "ImapMailBase"] diff --git a/gmailsorter/imap/__main__.py b/gmailsorter/imap/__main__.py new file mode 100644 index 0000000..96bbac1 --- /dev/null +++ b/gmailsorter/imap/__main__.py @@ -0,0 +1,99 @@ +import argparse +import os + +from gmailsorter import Imap + + +def command_line_parser(): + """ + Main function primarily used for the command line interface of the IMAP backend + """ + parser = argparse.ArgumentParser(prog="gmailsorter-imap") + parser.add_argument( + "--host", + help="IMAP server hostname e.g. imap.example.com .", + ) + parser.add_argument( + "--port", + type=int, + default=993, + help="IMAP server port - default: 993 .", + ) + parser.add_argument( + "--username", + help="IMAP account username.", + ) + parser.add_argument( + "--password-env", + default="IMAP_PASSWORD", + help=( + "Name of the environment variable holding the IMAP account password - " + "default: IMAP_PASSWORD ." + ), + ) + parser.add_argument( + "--no-ssl", + action="store_true", + help="Connect without SSL (IMAP4 instead of IMAP4_SSL).", + ) + parser.add_argument( + "-d", + "--database", + help="Connection string to connect to database e.g. sqlite:///email.db .", + ) + parser.add_argument( + "-u", + "--update", + action="store_true", + help="Update local database and retrain machine learning model.", + ) + parser.add_argument( + "-i", + "--identification", + help="User ID of the database user e.g. 1 .", + ) + parser.add_argument( + "-l", + "--label", + help="Email label (IMAP folder) to be filtered with machine learning.", + ) + args = parser.parse_args() + db_user_id = int(args.identification) if args.identification else 1 + password = os.environ.get(args.password_env) + if not args.host or not args.username: + print("Please provide --host and --username.") + elif not password: + print( + f"Please set the {args.password_env} environment variable to your IMAP password." + ) + else: + database = args.database or "sqlite:///email.db" + imap = Imap( + host=args.host, + port=args.port, + username=args.username, + password=password, + connection_str=database, + db_user_id=db_user_id, + use_ssl=not args.no_ssl, + email_download_format="metadata", + ) + if args.update: + imap.update_database(quick=False) + imap.fit_machine_learning_model_to_database( + n_estimators=100, + max_features=400, + random_state=42, + bootstrap=True, + include_deleted=False, + ) + elif args.label: + imap.filter_messages_from_server( + label=args.label, recommendation_ratio=0.9, label_prefix="labels_" + ) + else: + parser.print_help() + + +if __name__ == "__main__": + command_line_parser() diff --git a/gmailsorter/imap/authentication.py b/gmailsorter/imap/authentication.py new file mode 100644 index 0000000..bfa3200 --- /dev/null +++ b/gmailsorter/imap/authentication.py @@ -0,0 +1,21 @@ +from imaplib import IMAP4, IMAP4_SSL + + +def create_service(host, port, username, password, use_ssl=True): + """ + Open and log in to an IMAP connection. + + Args: + host (str): IMAP server hostname + port (int): IMAP server port + username (str): IMAP account username + password (str): IMAP account password + use_ssl (bool): connect via IMAP4_SSL (default) or plain IMAP4 + + Returns: + imaplib.IMAP4: logged-in IMAP connection + """ + connection_cls = IMAP4_SSL if use_ssl else IMAP4 + connection = connection_cls(host, port) + connection.login(username, password) + return connection diff --git a/gmailsorter/imap/mail.py b/gmailsorter/imap/mail.py new file mode 100644 index 0000000..15b1842 --- /dev/null +++ b/gmailsorter/imap/mail.py @@ -0,0 +1,355 @@ +import contextlib +import email +import imaplib +import re + +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from gmailsorter.base import get_email_database +from gmailsorter.base.mail import AbstractMailBox +from gmailsorter.imap.message import get_email_dict +from gmailsorter.ml import get_machine_learning_database + +_LIST_ENTRY_PATTERN = re.compile( + r'\((?P[^)]*)\)\s+(?:"(?P[^"]*)"|NIL)\s*(?P.*)' +) + +# Folder attributes which mark a folder as not usable as a sorting target: \Noselect +# (RFC 3501) plus the special-use attributes of RFC 6154. Without this the machine +# learning model can be trained on - and therefore recommend moving mail into - +# Trash, Spam, Sent or Drafts. +_SKIP_FOLDER_ATTRIBUTES = frozenset( + { + "\\noselect", + "\\all", + "\\archive", + "\\drafts", + "\\flagged", + "\\junk", + "\\sent", + "\\trash", + } +) + +# Many servers (GreenMail among them) do not advertise the RFC 6154 special-use +# attributes at all, so a deliberately short list of unambiguous special folder names +# is used as a fallback. Only exact (case-insensitive) matches are excluded, so a +# custom sorting folder such as "Sorted" or "MailSortInbox" is unaffected. +_SKIP_FOLDER_NAMES = frozenset( + { + "all mail", + "archive", + "deleted items", + "deleted messages", + "drafts", + "junk", + "junk e-mail", + "junk email", + "sent", + "sent items", + "sent messages", + "spam", + "trash", + } +) + + +def _decode_imap_bytes(value): + """ + Decode a bytes/str fragment of an IMAP response, returning None for anything else. + + Mailbox names are usually pure ASCII (modified UTF-7), but servers may send raw + UTF-8 literals - latin-1 is used as a lossless fallback so that a surprising + encoding never raises out of the LIST parser. + """ + if isinstance(value, str): + return value + if not isinstance(value, bytes): + return None + try: + return value.decode() + except UnicodeDecodeError: + return value.decode("latin-1") + + +class ImapMailBase(AbstractMailBox): + def close(self): + """ + Log out and close the IMAP connection. + + Errors raised while logging out are ignored on purpose - the connection is + being discarded anyway, and a connection which the server already dropped + must not turn closing it into a failure. + """ + with contextlib.suppress(imaplib.IMAP4.error, OSError): + self._service.logout() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + return False + + def _reconnect(self): + """ + Re-establish the IMAP connection and store it in self._service. + + Only a subclass which knows the connection details can do this, so the base + class signals that no reconnect is possible and the original connection error + is re-raised to the caller. + """ + raise NotImplementedError( + "ImapMailBase cannot reconnect because it does not know the connection " + "details - use gmailsorter.Imap for automatic reconnects." + ) + + def _run_with_reconnect(self, operation): + """ + Run an IMAP operation, retrying it exactly once on a dropped connection. + + IMAP servers commonly drop idle connections after 20-30 minutes and + update_database() can keep a single connection busy for far longer than that + on a large mailbox. imaplib signals a dropped connection with IMAP4.abort, + which would otherwise leave this instance permanently unusable. The operation + is re-run from the start after reconnecting, so a modification the server had + already applied before dropping the connection may report a failure on the + retry rather than being applied twice. + + Args: + operation (callable): zero-argument callable performing the IMAP calls + + Returns: + the return value of `operation` + """ + try: + return operation() + except imaplib.IMAP4.abort as error: + try: + self._reconnect() + except NotImplementedError: + raise error from None + return operation() + + def _get_label_translate_dict(self): + return self._run_with_reconnect(self._get_label_translate_dict_impl) + + def _get_label_translate_dict_impl(self): + status, mailbox_lst = self._service.list() + if status != "OK" or not mailbox_lst: + return {} + label_dict = {} + for entry in mailbox_lst: + parsed_entry = self._parse_list_entry(entry) + if parsed_entry is None: + continue + flags, delimiter, name = parsed_entry + if self._is_special_folder(flags=flags, delimiter=delimiter, name=name): + continue + label_dict[name] = name + return label_dict + + @staticmethod + def _is_special_folder(flags, delimiter, name): + """ + Check whether an IMAP folder is a special-purpose folder rather than a folder + emails may be sorted into. + + Args: + flags (list): folder attributes from the LIST response + delimiter (str/None): hierarchy delimiter from the LIST response + name (str): full folder name + + Returns: + bool: True if the folder must not be offered as a sorting label + """ + if any(flag.lower() in _SKIP_FOLDER_ATTRIBUTES for flag in flags): + return True + name_lst = [name.strip().lower()] + if delimiter: + # also check the leaf of a nested name such as "[Gmail]/Trash" + name_lst.append(name.rsplit(delimiter, 1)[-1].strip().lower()) + return any(candidate in _SKIP_FOLDER_NAMES for candidate in name_lst) + + def _search_email_on_server( + self, query_string="", label_lst=None, only_message_ids=False + ): + """ + Search emails either by a specific query or optionally limit your search to a list of labels + + Args: + query_string (str): not supported yet - must be empty + label_lst (list): list of IMAP folders to search; if empty, every folder is searched + only_message_ids (bool): return only the composite email IDs - default: false + + Returns: + list: list of composite "{folder}\\x1f{uid}" ids matching the search + """ + return self._run_with_reconnect( + lambda: self._search_email_on_server_impl( + query_string=query_string, + label_lst=label_lst, + only_message_ids=only_message_ids, + ) + ) + + def _search_email_on_server_impl( + self, query_string="", label_lst=None, only_message_ids=False + ): + if query_string: + raise NotImplementedError( + "Custom IMAP search queries are not supported yet, only label_lst filtering." + ) + if label_lst is None: + label_lst = [] + folder_lst = label_lst if len(label_lst) > 0 else list(self._label_dict.keys()) + message_id_lst = [ + f"{folder}\x1f{uid}" + for folder in folder_lst + for uid in self._search_folder(folder=folder) + ] + if only_message_ids: + return message_id_lst + else: + return [{"id": message_id} for message_id in message_id_lst] + + def _search_folder(self, folder): + status, _ = self._service.select(f'"{folder}"') + if status != "OK": + return [] + status, data = self._service.uid("search", None, "ALL") + if status != "OK" or data[0] is None: + return [] + return [ + uid.decode() if isinstance(uid, bytes) else uid for uid in data[0].split() + ] + + def _get_message_detail(self, message_id, email_format=None, metadata_headers=None): + """ + Fetch the raw message for a composite "{folder}\\x1f{uid}" id. + + BODY.PEEK[] is used rather than RFC822/BODY[] because the latter implicitly + set the \\Seen flag (RFC 3501), which would mark the whole mailbox as read + on every update_database() run. + + Returns: + tuple: (folder, uid, email.message.Message) + """ + return self._run_with_reconnect( + lambda: self._get_message_detail_impl(message_id=message_id) + ) + + def _get_message_detail_impl(self, message_id): + folder, uid = message_id.split("\x1f", 1) + status, _ = self._service.select(f'"{folder}"') + if status != "OK": + raise RuntimeError(f"Could not select IMAP folder {folder!r}") + status, data = self._service.uid("fetch", uid, "(BODY.PEEK[])") + if status != "OK" or not data or data[0] is None: + raise RuntimeError(f"Could not fetch IMAP message {message_id!r}") + raw_message = data[0][1] + parsed_message = email.message_from_bytes(raw_message) + return folder, uid, parsed_message + + def _modify_message_labels( + self, message_id, label_id_remove_lst=None, label_id_add_lst=None + ): + return self._run_with_reconnect( + lambda: self._modify_message_labels_impl( + message_id=message_id, + label_id_remove_lst=label_id_remove_lst, + label_id_add_lst=label_id_add_lst, + ) + ) + + def _modify_message_labels_impl( + self, message_id, label_id_remove_lst=None, label_id_add_lst=None + ): + if not label_id_add_lst: + return + folder, uid = message_id.split("\x1f", 1) + target_folder = label_id_add_lst[0] + status, _ = self._service.select(f'"{folder}"') + if status != "OK": + raise RuntimeError(f"Could not select IMAP folder {folder!r}") + if "MOVE" in self._service.capabilities: + status, _ = self._service.uid("move", uid, f'"{target_folder}"') + if status != "OK": + raise RuntimeError( + f"Could not move IMAP message {message_id!r} to {target_folder!r}" + ) + else: + status, _ = self._service.uid("copy", uid, f'"{target_folder}"') + if status != "OK": + raise RuntimeError( + f"Could not copy IMAP message {message_id!r} to {target_folder!r}" + ) + self._service.uid("store", uid, "+FLAGS", r"(\Deleted)") + if "UIDPLUS" in self._service.capabilities: + # RFC 4315 UID EXPUNGE - expunges only the message just copied + self._service.uid("expunge", uid) + else: + # Without UIDPLUS a bare EXPUNGE is the only option, and it also + # permanently removes any other message in this folder which is + # already flagged as \Deleted - an unavoidable limitation of + # servers supporting neither MOVE nor UIDPLUS. + self._service.expunge() + + def _get_labels_for_email(self, message_id): + folder, _uid = message_id.split("\x1f", 1) + return [folder] + + def _parse_message(self, message): + folder, uid, parsed_message = message + return get_email_dict(message=parsed_message, folder=folder, uid=uid) + + @staticmethod + def _parse_list_entry(entry): + """ + Parse a single entry of an IMAP LIST response. + + Handles the plain bytes/str form, the ``(header, literal_name)`` tuple form + imaplib returns when the server encodes the mailbox name as an IMAP literal, + and the unquoted ``NIL`` hierarchy delimiter which is legal per RFC 3501 for + servers without a folder hierarchy. + + Args: + entry (bytes/str/tuple/None): one element of ``imaplib.IMAP4.list()`` data + + Returns: + tuple/None: (flags, delimiter, name) or None if the entry is unparseable. + Unparseable entries are skipped rather than raised on, because + this runs from ``AbstractMailBox.__init__``. + """ + literal_name = None + if isinstance(entry, tuple): + try: + entry, literal_name = entry[0], _decode_imap_bytes(entry[1]) + except IndexError: + return None + if literal_name is None: + return None + decoded = _decode_imap_bytes(entry) + if decoded is None: + return None + match = _LIST_ENTRY_PATTERN.match(decoded) + if match is None: + return None + flags = match.group("flags").split() + delimiter = match.group("delimiter") + if literal_name is not None: + name = literal_name + else: + name = match.group("name").strip().strip('"') + if not name: + return None + return flags, delimiter, name + + @staticmethod + def _create_databases(connection_str): + engine = create_engine(connection_str) + session = sessionmaker(bind=engine)() + db_email = get_email_database(engine=engine, session=session) + db_ml = get_machine_learning_database(engine=engine, session=session) + return db_email, db_ml diff --git a/gmailsorter/imap/message.py b/gmailsorter/imap/message.py new file mode 100644 index 0000000..71ecc2b --- /dev/null +++ b/gmailsorter/imap/message.py @@ -0,0 +1,138 @@ +import email.header +import email.utils + +from gmailsorter.base.message import AbstractMessage, strip_html_tags + + +def get_email_dict(message, folder, uid): + try: + return Message(message=message, folder=folder, uid=uid).to_dict() + except (ValueError, KeyError) as e: + print(message, str(e)) + return None + + +class Message(AbstractMessage): + def __init__(self, message, folder, uid): + """ + Message class to parse a raw email.message.Message (as produced by + email.message_from_bytes() after an IMAP FETCH) into the common gmailsorter + email representation. + + Args: + message (email.message.Message): parsed RFC822 message + folder (str): IMAP mailbox/folder the message was fetched from + uid (str): IMAP UID of the message within `folder` + """ + self._message = message + self._folder = folder + self._uid = str(uid) + + def get_from(self): + from_header = self._message.get("From") + if from_header is None: + return None + addresses = [ + address for _, address in email.utils.getaddresses([from_header]) if address + ] + if len(addresses) == 1: + return addresses[0].lower() + return None + + def get_to(self): + return self._split_addresses(self._message.get_all("To")) + + def get_cc(self): + return self._split_addresses(self._message.get_all("Cc")) + + def get_label_ids(self): + return [self._folder] + + def get_subject(self): + return self._decode_header(self._message.get("Subject")) + + def get_date(self): + date_header = self._message.get("Date") + if date_header is None: + return None + return email.utils.parsedate_to_datetime(date_header) + + def get_content(self): + text_plain, text_html = None, None + if self._message.is_multipart(): + for part in self._message.walk(): + if part.get_content_maintype() == "multipart": + continue + if part.get_content_type() == "text/plain" and text_plain is None: + text_plain = self._decode_part(part) + elif part.get_content_type() == "text/html" and text_html is None: + text_html = self._decode_part(part) + elif self._message.get_content_type() == "text/plain": + text_plain = self._decode_part(self._message) + elif self._message.get_content_type() == "text/html": + text_html = self._decode_part(self._message) + if text_plain is not None: + return text_plain + elif text_html is not None: + return strip_html_tags(text_html) + else: + return None + + def get_thread_id(self): + references = self._message.get("References") + if references: + return references.split()[0] + in_reply_to = self._message.get("In-Reply-To") + if in_reply_to: + return in_reply_to.strip() + message_id = self._message.get("Message-ID") + if message_id: + return message_id.strip() + return self.get_email_id() + + def get_email_id(self): + return f"{self._folder}\x1f{self._uid}" + + @staticmethod + def _decode_header(header_value): + """ + Decode an email header into a plain str. + + Under the default compat32 policy, email.message.Message.get() normally + returns a str, but a header containing raw non-ASCII bytes that are not + valid RFC 2047 encoded-words (seen from some IMAP servers) is returned as + an email.header.Header instead. That object is not a str subclass, so it + fails SQLAlchemy's parameter binding - str() it first, then run RFC 2047 + decoding to resolve any encoded-words into text. + """ + if header_value is None: + return None + if isinstance(header_value, email.header.Header): + header_value = str(header_value) + decoded_chunks = email.header.decode_header(header_value) + return "".join( + ( + chunk.decode(charset or "utf-8", errors="replace") + if isinstance(chunk, bytes) + else chunk + ) + for chunk, charset in decoded_chunks + ) + + @staticmethod + def _decode_part(part): + payload = part.get_payload(decode=True) + if payload is None: + return "" + charset = part.get_content_charset() or "utf-8" + return payload.decode(charset, errors="replace") + + @staticmethod + def _split_addresses(header_values): + if not header_values: + return [] + return [ + address.lower() + for _, address in email.utils.getaddresses(header_values) + if address + ] diff --git a/gmailsorter/local.py b/gmailsorter/local.py index d074734..af43d12 100644 --- a/gmailsorter/local.py +++ b/gmailsorter/local.py @@ -2,6 +2,8 @@ from typing import Any from gmailsorter.google import GoogleMailBase, create_service +from gmailsorter.imap import ImapMailBase +from gmailsorter.imap import create_service as create_imap_service class Gmail(GoogleMailBase): @@ -65,3 +67,75 @@ def __init__( def load_client_secrets_file(client_secrets_file: str) -> dict[str, Any]: with open(client_secrets_file) as json_file: return json.load(json_file) + + +class Imap(ImapMailBase): + def __init__( + self, + host, + port, + username, + password, + connection_str, + db_user_id=1, + use_ssl=True, + email_download_format="metadata", + ): + """ + Imap class to manage Emails via a plain IMAP connection directly from Python + + The IMAP connection is kept open for the lifetime of the object. Call close() + when done, or use the object as a context manager: + + >>> with Imap(...) as imap: + ... imap.update_database() + + Args: + host (str): IMAP server hostname + port (int): IMAP server port, typically 993 for IMAP4_SSL or 143 for IMAP4 + username (str): IMAP account username + password (str): IMAP account password + connection_str (str): SQLalchemy compatible connection string to connect to the SQL database + db_user_id (int): Default 1 - set a user id when sharing a database with multiple users + use_ssl (bool): connect via IMAP4_SSL (default) or plain IMAP4 + email_download_format (str): unused for IMAP, kept for interface parity with Gmail + """ + self._connection_str = connection_str + # kept so the connection can be re-established when the server drops it + self._host = host + self._port = port + self._username = username + self._password = password + self._use_ssl = use_ssl + + database_email, database_ml = self._create_databases( + connection_str=self._connection_str + ) + + imap_connection = self._connect() + + super().__init__( + mail_service=imap_connection, + database_email=database_email, + database_ml=database_ml, + user_id=username, + db_user_id=db_user_id, + email_download_format=email_download_format, + ) + + def _connect(self): + return create_imap_service( + host=self._host, + port=self._port, + username=self._username, + password=self._password, + use_ssl=self._use_ssl, + ) + + def _reconnect(self): + """ + Re-establish the IMAP connection after the server dropped it, discarding the + old connection first. + """ + self.close() + self._service = self._connect() diff --git a/notebooks/demo.ipynb b/notebooks/demo.ipynb index 29d0a84..0e5767f 100644 --- a/notebooks/demo.ipynb +++ b/notebooks/demo.ipynb @@ -127,7 +127,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "907c830a", "metadata": {}, "outputs": [ @@ -141,7 +141,9 @@ } ], "source": [ - "gmail.filter_messages_from_server(label=\"9-emails-to-sort\", recommendation_ratio=0.9)" + "gmail.filter_messages_from_server(\n", + " label=\"9-emails-to-sort\", recommendation_ratio=0.9, label_prefix=\"labels_Label_\"\n", + ")" ] } ], diff --git a/pyproject.toml b/pyproject.toml index a58c406..34fa025 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ Homepage = "https://github.com/jan-janssen/gmailsorter" gmailsorter = "gmailsorter.__main__:command_line_parser" gmailsorter-daemon = "gmailsorter.daemon.__main__:command_line_parser" gmailsorter-app = "gmailsorter.webapp.app:run_app" +gmailsorter-imap = "gmailsorter.imap.__main__:command_line_parser" [tool.ruff.lint] select = [ diff --git a/tests/test_google_integration_units.py b/tests/test_google_integration_units.py index 1bf6013..1f1f5f0 100644 --- a/tests/test_google_integration_units.py +++ b/tests/test_google_integration_units.py @@ -361,8 +361,8 @@ def test_update_database_quick_and_full_paths(self): db_email.update_labels.assert_not_called() store_mock.assert_called_once_with(message_id_lst=["new2"], email_format=None) - @patch("gmailsorter.google.mail.get_predictions_from_machine_learning_models") - @patch("gmailsorter.google.mail.encode_df_for_machine_learning") + @patch("gmailsorter.base.mail.get_predictions_from_machine_learning_models") + @patch("gmailsorter.base.mail.encode_df_for_machine_learning") def test_filter_messages_from_server(self, encode_mock, predict_mock): service = self._create_mock_service_with_labels() db_ml = MagicMock() @@ -395,8 +395,8 @@ def test_filter_messages_from_server(self, encode_mock, predict_mock): mail.filter_messages_from_server("Inbox") encode_mock.assert_not_called() - @patch("gmailsorter.google.mail.fit_machine_learning_models") - @patch("gmailsorter.google.mail.encode_df_for_machine_learning") + @patch("gmailsorter.base.mail.fit_machine_learning_models") + @patch("gmailsorter.base.mail.encode_df_for_machine_learning") def test_fit_machine_learning_model_to_database(self, encode_mock, fit_mock): service = self._create_mock_service_with_labels() db_ml = MagicMock() @@ -440,8 +440,8 @@ def test_fit_machine_learning_model_to_database(self, encode_mock, fit_mock): db_ml.store_models.assert_called_once() self.assertEqual(mail.get_all_emails_in_database().iloc[0]["id"], "x") - @patch("gmailsorter.google.mail.fit_machine_learning_models") - @patch("gmailsorter.google.mail.encode_df_for_machine_learning") + @patch("gmailsorter.base.mail.fit_machine_learning_models") + @patch("gmailsorter.base.mail.encode_df_for_machine_learning") def test_fit_machine_learning_model_to_database_forwards_max_workers( self, encode_mock, fit_mock ): diff --git a/tests/test_google_message.py b/tests/test_google_message.py index 934e331..89314b3 100644 --- a/tests/test_google_message.py +++ b/tests/test_google_message.py @@ -2,7 +2,8 @@ from unittest import TestCase from datetime import datetime from datetime import datetime, timezone, timedelta -from gmailsorter.google.message import Message, MLStripper, get_email_dict +from gmailsorter.base.message import strip_html_tags +from gmailsorter.google.message import Message, get_email_dict class MessageTest(TestCase): @@ -198,9 +199,9 @@ def test_get_content_missing_body_data_returns_empty_string(self): self.assertEqual(message.get_content(), "") def test_mlstripper_removes_tags(self): - stripper = MLStripper() - stripper.feed("
Hello World
") - self.assertEqual(stripper.get_data(), "Hello World") + self.assertEqual( + strip_html_tags("
Hello World
"), "Hello World" + ) def test_get_email_dict_catches_value_error_and_returns_none(self): message_dict = { diff --git a/tests/test_imap_cli.py b/tests/test_imap_cli.py new file mode 100644 index 0000000..04fcf69 --- /dev/null +++ b/tests/test_imap_cli.py @@ -0,0 +1,127 @@ +import os +from unittest import TestCase +from unittest.mock import patch + +from gmailsorter.imap.__main__ import command_line_parser + + +class ImapCliTest(TestCase): + @patch("gmailsorter.imap.__main__.Imap") + def test_update_wires_imap_and_triggers_update(self, imap_cls): + imap_instance = imap_cls.return_value + os.environ["IMAP_PASSWORD"] = "secret" + try: + with patch( + "sys.argv", + [ + "gmailsorter-imap", + "--host", + "localhost", + "--port", + "993", + "--username", + "user", + "-d", + "sqlite:///:memory:", + "-u", + ], + ): + command_line_parser() + finally: + del os.environ["IMAP_PASSWORD"] + + imap_cls.assert_called_once_with( + host="localhost", + port=993, + username="user", + password="secret", + connection_str="sqlite:///:memory:", + db_user_id=1, + use_ssl=True, + email_download_format="metadata", + ) + imap_instance.update_database.assert_called_once_with(quick=False) + imap_instance.fit_machine_learning_model_to_database.assert_called_once_with( + n_estimators=100, + max_features=400, + random_state=42, + bootstrap=True, + include_deleted=False, + ) + + @patch("gmailsorter.imap.__main__.Imap") + def test_label_wires_imap_and_triggers_filter(self, imap_cls): + imap_instance = imap_cls.return_value + os.environ["IMAP_PASSWORD"] = "secret" + try: + with patch( + "sys.argv", + [ + "gmailsorter-imap", + "--host", + "localhost", + "--username", + "user", + "-d", + "sqlite:///:memory:", + "-l", + "MailSortInbox", + ], + ): + command_line_parser() + finally: + del os.environ["IMAP_PASSWORD"] + + imap_instance.filter_messages_from_server.assert_called_once_with( + label="MailSortInbox", + recommendation_ratio=0.9, + label_prefix="labels_", + ) + + @patch("gmailsorter.imap.__main__.Imap") + def test_no_update_or_label_prints_help(self, imap_cls): + imap_instance = imap_cls.return_value + os.environ["IMAP_PASSWORD"] = "secret" + try: + with patch( + "sys.argv", + [ + "gmailsorter-imap", + "--host", + "localhost", + "--username", + "user", + "-d", + "sqlite:///:memory:", + ], + ): + command_line_parser() + finally: + del os.environ["IMAP_PASSWORD"] + + imap_instance.update_database.assert_not_called() + imap_instance.filter_messages_from_server.assert_not_called() + + @patch("gmailsorter.imap.__main__.Imap") + def test_missing_password_env_skips_wiring(self, imap_cls): + os.environ.pop("IMAP_PASSWORD", None) + with patch( + "sys.argv", + ["gmailsorter-imap", "--host", "localhost", "--username", "user"], + ): + command_line_parser() + + imap_cls.assert_not_called() + + @patch("gmailsorter.imap.__main__.Imap") + def test_missing_host_skips_wiring(self, imap_cls): + with patch("sys.argv", ["gmailsorter-imap", "--username", "user"]): + command_line_parser() + + imap_cls.assert_not_called() + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/tests/test_imap_integration_units.py b/tests/test_imap_integration_units.py new file mode 100644 index 0000000..4aaa68b --- /dev/null +++ b/tests/test_imap_integration_units.py @@ -0,0 +1,595 @@ +import imaplib +from unittest import TestCase +from unittest.mock import MagicMock, patch + +from gmailsorter.imap.authentication import create_service +from gmailsorter.imap.mail import ImapMailBase +from gmailsorter.local import Imap + + +class ReconnectingImapMailBase(ImapMailBase): + """ + Test double which reconnects by swapping in the next prepared mock connection. + + This stands in for gmailsorter.local.Imap, which is the class that actually knows + the connection details, without requiring a real IMAP server. + """ + + def __init__(self, service_lst, **kwargs): + self.reconnect_count = 0 + self._service_lst = list(service_lst) + super().__init__(mail_service=self._service_lst.pop(0), **kwargs) + + def _reconnect(self): + self.reconnect_count += 1 + self._service = self._service_lst.pop(0) + + +class TestImapAuthentication(TestCase): + @patch("gmailsorter.imap.authentication.IMAP4_SSL") + def test_create_service_uses_ssl_by_default(self, imap_ssl_cls): + connection = imap_ssl_cls.return_value + + result = create_service( + host="localhost", port=993, username="user", password="secret" + ) + + imap_ssl_cls.assert_called_once_with("localhost", 993) + connection.login.assert_called_once_with("user", "secret") + self.assertIs(result, connection) + + @patch("gmailsorter.imap.authentication.IMAP4") + def test_create_service_without_ssl(self, imap_cls): + connection = imap_cls.return_value + + result = create_service( + host="localhost", + port=143, + username="user", + password="secret", + use_ssl=False, + ) + + imap_cls.assert_called_once_with("localhost", 143) + connection.login.assert_called_once_with("user", "secret") + self.assertIs(result, connection) + + +class TestImapMailBase(TestCase): + def _create_mock_service_with_folders(self, folders=None): + service = MagicMock() + service.capabilities = ["IMAP4rev1", "MOVE"] + if folders is None: + folders = [ + b'(\\HasNoChildren) "/" "INBOX"', + b'(\\HasNoChildren) "/" "MailSortInbox"', + b'(\\Noselect \\HasChildren) "/" "[Gmail]"', + ] + service.list.return_value = ("OK", folders) + return service + + def test_get_label_translate_dict_skips_noselect(self): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + + self.assertEqual(sorted(mail.labels), ["INBOX", "MailSortInbox"]) + + def test_get_label_translate_dict_skips_special_use_attributes(self): + service = self._create_mock_service_with_folders( + folders=[ + b'(\\HasNoChildren \\Trash) "/" "Papierkorb"', + b'(\\HasNoChildren \\Junk) "/" "Unerwuenscht"', + b'(\\HasNoChildren \\Sent) "/" "Gesendet"', + b'(\\HasNoChildren \\Drafts) "/" "Entwuerfe"', + b'(\\HasNoChildren \\Archive) "/" "Ablage"', + b'(\\HasNoChildren \\All) "/" "Alle"', + b'(\\HasNoChildren \\Flagged) "/" "Markiert"', + b'(\\HasNoChildren) "/" "MailSortInbox"', + ] + ) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, ["MailSortInbox"]) + + def test_get_label_translate_dict_skips_special_names_without_attributes(self): + service = self._create_mock_service_with_folders( + folders=[ + b'(\\HasNoChildren) "/" "Trash"', + b'(\\HasNoChildren) "/" "junk e-mail"', + b'(\\HasNoChildren) "/" "Deleted Items"', + b'(\\HasNoChildren) "/" "SPAM"', + b'(\\HasNoChildren) "/" "Sent Items"', + b'(\\HasNoChildren) "/" "Drafts"', + b'(\\HasNoChildren) "/" "All Mail"', + b'(\\HasNoChildren) "/" "[Gmail]/Trash"', + b'(\\HasNoChildren) "/" "MailSortInbox"', + ] + ) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, ["MailSortInbox"]) + + def test_get_label_translate_dict_keeps_custom_folders(self): + service = self._create_mock_service_with_folders( + folders=[ + b'(\\HasNoChildren) "/" "MailSortInbox"', + b'(\\HasNoChildren) "/" "Sorted"', + b'(\\HasNoChildren) "/" "Archived Projects"', + b'(\\HasNoChildren) "/" "Trashcan Design"', + b'(\\HasNoChildren) "/" "INBOX"', + ] + ) + mail = ImapMailBase(mail_service=service) + + self.assertEqual( + sorted(mail.labels), + [ + "Archived Projects", + "INBOX", + "MailSortInbox", + "Sorted", + "Trashcan Design", + ], + ) + + def test_get_label_translate_dict_accepts_nil_delimiter(self): + service = self._create_mock_service_with_folders( + folders=[b"(\\HasNoChildren) NIL INBOX"] + ) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, ["INBOX"]) + + def test_get_label_translate_dict_accepts_literal_name_tuple(self): + service = self._create_mock_service_with_folders( + folders=[(b'(\\HasNoChildren) "/" {11}', b"MailSortBox"), b")"] + ) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, ["MailSortBox"]) + + def test_get_label_translate_dict_handles_empty_mailbox_list(self): + service = self._create_mock_service_with_folders(folders=[None]) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, []) + + def test_get_label_translate_dict_returns_empty_on_list_failure(self): + service = self._create_mock_service_with_folders() + service.list.return_value = ("NO", None) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, []) + + def test_get_label_translate_dict_skips_unparseable_entries(self): + service = self._create_mock_service_with_folders( + folders=[b"total garbage", b'(\\HasNoChildren) "/" "MailSortInbox"'] + ) + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail.labels, ["MailSortInbox"]) + + def test_parse_list_entry_returns_none_for_unparseable_input(self): + self.assertIsNone(ImapMailBase._parse_list_entry(None)) + self.assertIsNone(ImapMailBase._parse_list_entry(b"not a list response")) + self.assertIsNone(ImapMailBase._parse_list_entry((b'(\\Noselect) "/" {3}',))) + self.assertIsNone(ImapMailBase._parse_list_entry(b'(\\HasNoChildren) "/" ')) + + def test_parse_list_entry_returns_none_when_literal_name_undecodable(self): + self.assertIsNone( + ImapMailBase._parse_list_entry((b'(\\HasNoChildren) "/" {3}', None)) + ) + + def test_parse_list_entry_accepts_plain_str_entry(self): + result = ImapMailBase._parse_list_entry('(\\HasNoChildren) "/" "INBOX"') + + self.assertEqual(result, (["\\HasNoChildren"], "/", "INBOX")) + + def test_parse_list_entry_falls_back_to_latin1_on_invalid_utf8(self): + result = ImapMailBase._parse_list_entry(b'(\\HasNoChildren) "/" "Caf\xe9"') + + self.assertEqual(result, (["\\HasNoChildren"], "/", "Caf\xe9")) + + def test_search_email_on_server_single_folder(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [b"1 2"]) + mail = ImapMailBase(mail_service=service) + + ids = mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + service.select.assert_called_once_with('"INBOX"') + service.uid.assert_called_once_with("search", None, "ALL") + self.assertEqual(ids, ["INBOX\x1f1", "INBOX\x1f2"]) + + def test_search_email_on_server_all_folders_when_no_label(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [b"5"]) + mail = ImapMailBase(mail_service=service) + + ids = mail._search_email_on_server(only_message_ids=True) + + self.assertEqual( + service.select.call_args_list, + [(('"INBOX"',),), (('"MailSortInbox"',),)], + ) + self.assertEqual(ids, ["INBOX\x1f5", "MailSortInbox\x1f5"]) + + def test_search_email_on_server_rejects_query_string(self): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + + with self.assertRaises(NotImplementedError): + mail._search_email_on_server(query_string="SUBJECT foo") + + def test_search_email_on_server_returns_dicts_by_default(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [b"1 2"]) + mail = ImapMailBase(mail_service=service) + + result = mail._search_email_on_server(label_lst=["INBOX"]) + + self.assertEqual(result, [{"id": "INBOX\x1f1"}, {"id": "INBOX\x1f2"}]) + + def test_search_folder_returns_empty_list_when_select_fails(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("NO", [b"failed"]) + mail = ImapMailBase(mail_service=service) + + ids = mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + service.uid.assert_not_called() + self.assertEqual(ids, []) + + def test_search_folder_returns_empty_list_when_search_fails(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("NO", [None]) + mail = ImapMailBase(mail_service=service) + + ids = mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + self.assertEqual(ids, []) + + def test_search_folder_returns_empty_list_when_search_data_is_none(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [None]) + mail = ImapMailBase(mail_service=service) + + ids = mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + self.assertEqual(ids, []) + + def test_get_message_detail_selects_and_fetches(self): + service = self._create_mock_service_with_folders() + raw_message = b"Subject: hi\r\nFrom: a@b.com\r\nTo: c@d.com\r\n\r\nbody" + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [(b"1 (BODY[] {10}", raw_message)]) + mail = ImapMailBase(mail_service=service) + + folder, uid, message = mail._get_message_detail(message_id="INBOX\x1f7") + + service.select.assert_called_once_with('"INBOX"') + service.uid.assert_called_once_with("fetch", "7", "(BODY.PEEK[])") + self.assertEqual(folder, "INBOX") + self.assertEqual(uid, "7") + self.assertEqual(message["Subject"], "hi") + + def test_get_message_detail_raises_when_select_fails(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("NO", [b"failed"]) + mail = ImapMailBase(mail_service=service) + + with self.assertRaises(RuntimeError): + mail._get_message_detail(message_id="INBOX\x1f7") + + def test_get_message_detail_raises_when_fetch_fails(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("NO", [None]) + mail = ImapMailBase(mail_service=service) + + with self.assertRaises(RuntimeError): + mail._get_message_detail(message_id="INBOX\x1f7") + + def test_get_labels_for_email_from_composite_id(self): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + + self.assertEqual(mail._get_labels_for_email("INBOX\x1f7"), ["INBOX"]) + + def test_modify_message_labels_uses_move_when_supported(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [b"1"]) + mail = ImapMailBase(mail_service=service) + + mail._modify_message_labels( + message_id="INBOX\x1f7", + label_id_remove_lst=["INBOX"], + label_id_add_lst=["MailSortInbox"], + ) + + service.select.assert_called_once_with('"INBOX"') + service.uid.assert_called_once_with("move", "7", '"MailSortInbox"') + service.expunge.assert_not_called() + + def test_modify_message_labels_falls_back_to_copy_delete(self): + service = self._create_mock_service_with_folders() + service.capabilities = ["IMAP4rev1"] + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [b"1"]) + mail = ImapMailBase(mail_service=service) + + mail._modify_message_labels( + message_id="INBOX\x1f7", + label_id_remove_lst=["INBOX"], + label_id_add_lst=["MailSortInbox"], + ) + + self.assertEqual( + service.uid.call_args_list, + [ + (("copy", "7", '"MailSortInbox"'),), + (("store", "7", "+FLAGS", r"(\Deleted)"),), + ], + ) + service.expunge.assert_called_once() + + def test_modify_message_labels_uses_uid_expunge_with_uidplus(self): + service = self._create_mock_service_with_folders() + service.capabilities = ["IMAP4rev1", "UIDPLUS"] + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("OK", [b"1"]) + mail = ImapMailBase(mail_service=service) + + mail._modify_message_labels( + message_id="INBOX\x1f7", + label_id_remove_lst=["INBOX"], + label_id_add_lst=["MailSortInbox"], + ) + + self.assertEqual( + service.uid.call_args_list, + [ + (("copy", "7", '"MailSortInbox"'),), + (("store", "7", "+FLAGS", r"(\Deleted)"),), + (("expunge", "7"),), + ], + ) + service.expunge.assert_not_called() + + def test_modify_message_labels_raises_when_select_fails(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("NO", [b"failed"]) + mail = ImapMailBase(mail_service=service) + + with self.assertRaises(RuntimeError): + mail._modify_message_labels( + message_id="INBOX\x1f7", label_id_add_lst=["MailSortInbox"] + ) + + def test_modify_message_labels_raises_when_move_fails(self): + service = self._create_mock_service_with_folders() + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("NO", [None]) + mail = ImapMailBase(mail_service=service) + + with self.assertRaises(RuntimeError): + mail._modify_message_labels( + message_id="INBOX\x1f7", label_id_add_lst=["MailSortInbox"] + ) + + def test_modify_message_labels_raises_when_copy_fails(self): + service = self._create_mock_service_with_folders() + service.capabilities = ["IMAP4rev1"] + service.select.return_value = ("OK", [b"1"]) + service.uid.return_value = ("NO", [None]) + mail = ImapMailBase(mail_service=service) + + with self.assertRaises(RuntimeError): + mail._modify_message_labels( + message_id="INBOX\x1f7", label_id_add_lst=["MailSortInbox"] + ) + + def test_modify_message_labels_noop_without_target(self): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + + mail._modify_message_labels(message_id="INBOX\x1f7") + + service.select.assert_not_called() + + @patch("gmailsorter.imap.mail.get_email_dict") + def test_parse_message_delegates_to_get_email_dict(self, get_email_dict_mock): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + get_email_dict_mock.return_value = {"id": "INBOX\x1f7"} + + result = mail._parse_message(("INBOX", "7", "raw")) + + get_email_dict_mock.assert_called_once_with( + message="raw", folder="INBOX", uid="7" + ) + self.assertEqual(result, {"id": "INBOX\x1f7"}) + + def test_create_databases(self): + with ( + patch("gmailsorter.imap.mail.create_engine") as create_engine_mock, + patch("gmailsorter.imap.mail.sessionmaker") as sessionmaker_mock, + patch("gmailsorter.imap.mail.get_email_database") as get_email_db_mock, + patch( + "gmailsorter.imap.mail.get_machine_learning_database" + ) as get_ml_db_mock, + ): + engine = MagicMock() + session = MagicMock() + create_engine_mock.return_value = engine + sessionmaker_mock.return_value.return_value = session + get_email_db_mock.return_value = "EMAIL_DB" + get_ml_db_mock.return_value = "ML_DB" + + dbs = ImapMailBase._create_databases("sqlite:///file.db") + + self.assertEqual(dbs, ("EMAIL_DB", "ML_DB")) + + def test_close_logs_out(self): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + + mail.close() + + service.logout.assert_called_once_with() + + def test_close_ignores_logout_failure(self): + for error in (imaplib.IMAP4.abort("connection lost"), OSError("socket gone")): + with self.subTest(error=type(error).__name__): + service = self._create_mock_service_with_folders() + service.logout.side_effect = error + mail = ImapMailBase(mail_service=service) + + mail.close() + + service.logout.assert_called_once_with() + + def test_context_manager_closes_connection(self): + service = self._create_mock_service_with_folders() + + with ImapMailBase(mail_service=service) as mail: + self.assertEqual(sorted(mail.labels), ["INBOX", "MailSortInbox"]) + service.logout.assert_not_called() + + service.logout.assert_called_once_with() + + def test_abort_propagates_without_reconnect_support(self): + service = self._create_mock_service_with_folders() + mail = ImapMailBase(mail_service=service) + service.select.side_effect = imaplib.IMAP4.abort("connection lost") + + with self.assertRaises(imaplib.IMAP4.abort): + mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + def test_search_reconnects_and_retries_once_after_abort(self): + dead_service = self._create_mock_service_with_folders() + dead_service.select.side_effect = imaplib.IMAP4.abort("connection lost") + fresh_service = self._create_mock_service_with_folders() + fresh_service.select.return_value = ("OK", [b"1"]) + fresh_service.uid.return_value = ("OK", [b"3"]) + mail = ReconnectingImapMailBase(service_lst=[dead_service, fresh_service]) + + ids = mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + self.assertEqual(mail.reconnect_count, 1) + self.assertEqual(ids, ["INBOX\x1f3"]) + + def test_get_message_detail_reconnects_and_retries_once_after_abort(self): + raw_message = b"Subject: hi\r\nFrom: a@b.com\r\nTo: c@d.com\r\n\r\nbody" + dead_service = self._create_mock_service_with_folders() + dead_service.select.side_effect = imaplib.IMAP4.abort("connection lost") + fresh_service = self._create_mock_service_with_folders() + fresh_service.select.return_value = ("OK", [b"1"]) + fresh_service.uid.return_value = ("OK", [(b"1 (BODY[] {10}", raw_message)]) + mail = ReconnectingImapMailBase(service_lst=[dead_service, fresh_service]) + + folder, uid, message = mail._get_message_detail(message_id="INBOX\x1f7") + + self.assertEqual(mail.reconnect_count, 1) + self.assertEqual((folder, uid), ("INBOX", "7")) + self.assertEqual(message["Subject"], "hi") + + def test_retry_is_attempted_only_once(self): + dead_service = self._create_mock_service_with_folders() + dead_service.select.side_effect = imaplib.IMAP4.abort("connection lost") + still_dead_service = self._create_mock_service_with_folders() + still_dead_service.select.side_effect = imaplib.IMAP4.abort("connection lost") + mail = ReconnectingImapMailBase( + service_lst=[dead_service, still_dead_service], + ) + + with self.assertRaises(imaplib.IMAP4.abort): + mail._search_email_on_server(label_lst=["INBOX"], only_message_ids=True) + + self.assertEqual(mail.reconnect_count, 1) + + +class TestImapLocalHelpers(TestCase): + @patch("gmailsorter.local.ImapMailBase.__init__", return_value=None) + @patch("gmailsorter.local.create_imap_service") + @patch("gmailsorter.local.Imap._create_databases") + def test_imap_initialization_wiring( + self, create_databases_mock, create_service_mock, base_init_mock + ): + db_email, db_ml = MagicMock(), MagicMock() + create_databases_mock.return_value = (db_email, db_ml) + connection = MagicMock() + create_service_mock.return_value = connection + + Imap( + host="localhost", + port=993, + username="user", + password="secret", + connection_str="sqlite:///:memory:", + db_user_id=4, + ) + + create_databases_mock.assert_called_once_with( + connection_str="sqlite:///:memory:" + ) + create_service_mock.assert_called_once_with( + host="localhost", + port=993, + username="user", + password="secret", + use_ssl=True, + ) + base_init_mock.assert_called_once_with( + mail_service=connection, + database_email=db_email, + database_ml=db_ml, + user_id="user", + db_user_id=4, + email_download_format="metadata", + ) + + @patch("gmailsorter.local.ImapMailBase.__init__", return_value=None) + @patch("gmailsorter.local.create_imap_service") + @patch("gmailsorter.local.Imap._create_databases") + def test_imap_reconnect_replaces_the_connection( + self, create_databases_mock, create_service_mock, base_init_mock + ): + create_databases_mock.return_value = (MagicMock(), MagicMock()) + dead_connection, fresh_connection = MagicMock(), MagicMock() + create_service_mock.side_effect = [dead_connection, fresh_connection] + + imap = Imap( + host="mail.example.test", + port=143, + username="user", + password="secret", + connection_str="sqlite:///:memory:", + use_ssl=False, + ) + # ImapMailBase.__init__ is mocked out above, so _service is set by hand here + imap._service = dead_connection + + imap._reconnect() + + dead_connection.logout.assert_called_once_with() + self.assertIs(imap._service, fresh_connection) + create_service_mock.assert_called_with( + host="mail.example.test", + port=143, + username="user", + password="secret", + use_ssl=False, + ) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/tests/test_imap_message.py b/tests/test_imap_message.py new file mode 100644 index 0000000..ccac012 --- /dev/null +++ b/tests/test_imap_message.py @@ -0,0 +1,173 @@ +from datetime import datetime +from email.header import Header +from email.message import EmailMessage, Message as EmailLibMessage +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText +from unittest import TestCase + +from gmailsorter.imap.message import Message, get_email_dict + + +class MessageTest(TestCase): + @classmethod + def setUpClass(cls) -> None: + msg = EmailMessage() + msg["Subject"] = "Test Email Subject" + msg["From"] = "sender@server.net" + msg["To"] = "me@mail.com, friend@provider.org" + msg["Date"] = "Fri, 11 Feb 2022 18:08:46 +0100" + msg["Message-ID"] = "" + msg.set_content("Hello world") + cls._message = msg + cls.message = Message(message=msg, folder="INBOX", uid="42") + + def test_subject(self): + self.assertEqual(self.message.get_subject(), "Test Email Subject") + + def test_subject_encoded_word(self): + msg = EmailMessage() + msg["Subject"] = "Exclusieve Nieuwsbrief • Binobet" + message = Message(message=msg, folder="INBOX", uid="1") + self.assertEqual(message.get_subject(), "Exclusieve Nieuwsbrief • Binobet") + + def test_subject_header_object_is_coerced_to_str(self): + # Some servers/Python versions cause Message.get() to return an + # email.header.Header instance instead of a str, which used to crash + # SQLAlchemy's parameter binding when inserted into the database. + msg = EmailLibMessage() + msg["Subject"] = Header("Exclusieve Nieuwsbrief • Binobet", "utf-8") + message = Message(message=msg, folder="INBOX", uid="1") + subject = message.get_subject() + self.assertIsInstance(subject, str) + self.assertEqual(subject, "Exclusieve Nieuwsbrief • Binobet") + + def test_from(self): + self.assertEqual(self.message.get_from(), "sender@server.net") + + def test_to(self): + self.assertEqual(self.message.get_to(), ["me@mail.com", "friend@provider.org"]) + + def test_cc_empty(self): + self.assertEqual(self.message.get_cc(), []) + + def test_email_id(self): + self.assertEqual(self.message.get_email_id(), "INBOX\x1f42") + + def test_thread_id_falls_back_to_message_id(self): + self.assertEqual(self.message.get_thread_id(), "") + + def test_label_ids(self): + self.assertEqual(self.message.get_label_ids(), ["INBOX"]) + + def test_get_date(self): + self.assertEqual( + self.message.get_date(), + datetime.strptime( + "Fri, 11 Feb 2022 18:08:46 +0100", "%a, %d %b %Y %H:%M:%S %z" + ), + ) + + def test_get_content(self): + self.assertEqual(self.message.get_content().strip(), "Hello world") + + def test_get_content_html_fallback(self): + html_msg = EmailMessage() + html_msg["Subject"] = "HTML" + html_msg["From"] = "sender@server.net" + html_msg["To"] = "me@mail.com" + html_msg["Date"] = "Fri, 11 Feb 2022 18:08:46 +0100" + html_msg.set_content("

Hello World

", subtype="html") + message = Message(message=html_msg, folder="INBOX", uid="43") + + self.assertEqual(message.get_content().strip(), "Hello World") + + def test_thread_id_uses_references_header(self): + msg = EmailMessage() + msg["Subject"] = "Re: Test" + msg["References"] = " " + msg["Message-ID"] = "" + message = Message(message=msg, folder="INBOX", uid="44") + + self.assertEqual(message.get_thread_id(), "") + + def test_from_with_multiple_addresses_is_none(self): + msg = EmailMessage() + msg["From"] = "a@server.net, b@server.net" + message = Message(message=msg, folder="INBOX", uid="45") + + self.assertIsNone(message.get_from()) + + def test_get_from_missing_header_returns_none(self): + msg = EmailMessage() + message = Message(message=msg, folder="INBOX", uid="46") + + self.assertIsNone(message.get_from()) + + def test_get_date_missing_header_returns_none(self): + msg = EmailMessage() + message = Message(message=msg, folder="INBOX", uid="47") + + self.assertIsNone(message.get_date()) + + def test_get_content_multipart_prefers_plain_over_html(self): + outer = MIMEMultipart("mixed") + inner = MIMEMultipart("alternative") + inner.attach(MIMEText("Hello world", "plain")) + inner.attach(MIMEText("

Hello World

", "html")) + outer.attach(inner) + message = Message(message=outer, folder="INBOX", uid="48") + + self.assertEqual(message.get_content().strip(), "Hello world") + + def test_get_content_returns_none_for_unknown_mimetype(self): + msg = EmailMessage() + msg.set_content(b"\x00\x01", maintype="application", subtype="octet-stream") + message = Message(message=msg, folder="INBOX", uid="49") + + self.assertIsNone(message.get_content()) + + def test_thread_id_uses_in_reply_to_when_no_references(self): + msg = EmailMessage() + msg["In-Reply-To"] = " " + msg["Message-ID"] = "" + message = Message(message=msg, folder="INBOX", uid="50") + + self.assertEqual(message.get_thread_id(), "") + + def test_thread_id_falls_back_to_email_id_without_any_headers(self): + msg = EmailMessage() + message = Message(message=msg, folder="INBOX", uid="51") + + self.assertEqual(message.get_thread_id(), "INBOX\x1f51") + + def test_decode_part_returns_empty_string_when_payload_is_none(self): + multipart_msg = MIMEMultipart("mixed") + + self.assertEqual(Message._decode_part(multipart_msg), "") + + def test_get_email_dict_catches_value_error_and_returns_none(self): + msg = EmailMessage() + msg["Date"] = "Mon, 32 Jan 2024 25:99:99 +0000" + + self.assertIsNone(get_email_dict(msg, folder="INBOX", uid="52")) + + def test_get_email_dict(self): + result = get_email_dict(self._message, folder="INBOX", uid="42") + content = result.pop("content") + + self.assertEqual(content.strip(), "Hello world") + self.assertEqual( + result, + { + "cc": [], + "date": datetime.strptime( + "Fri, 11 Feb 2022 18:08:46 +0100", "%a, %d %b %Y %H:%M:%S %z" + ), + "from": "sender@server.net", + "id": "INBOX\x1f42", + "labels": ["INBOX"], + "subject": "Test Email Subject", + "threads": "", + "to": ["me@mail.com", "friend@provider.org"], + }, + ) diff --git a/tests/test_imap_service_integration.py b/tests/test_imap_service_integration.py new file mode 100644 index 0000000..2664597 --- /dev/null +++ b/tests/test_imap_service_integration.py @@ -0,0 +1,181 @@ +import os +import smtplib +import time +import unittest +import uuid +from email.message import EmailMessage +from imaplib import IMAP4 +from unittest.mock import patch + +from gmailsorter.local import Imap + +#: When this environment variable is set to any non-empty value, an unreachable IMAP +#: server is a test failure rather than a skip. CI sets it so the imap-integration job +#: cannot report green when the GreenMail service never came up; local runs leave it +#: unset and keep the clean skip. +IMAP_INTEGRATION_REQUIRED = "IMAP_INTEGRATION_REQUIRED" + + +def imap_integration_required(): + """ + Returns: + bool: True if a missing IMAP test server must fail instead of skip + """ + return bool(os.environ.get(IMAP_INTEGRATION_REQUIRED, "").strip()) + + +class TestImapServiceIntegration(unittest.TestCase): + smtp_host = os.environ.get("TEST_SMTP_HOST", "localhost") + smtp_port = int(os.environ.get("TEST_SMTP_PORT", "3025")) + imap_host = os.environ.get("TEST_IMAP_HOST", "localhost") + imap_port = int(os.environ.get("TEST_IMAP_PORT", "3143")) + username = os.environ.get("TEST_IMAP_USERNAME", "testuser") + recipient = os.environ.get("TEST_EMAIL", "testuser@example.test") + password = os.environ.get("TEST_EMAIL_PASSWORD", "secret") + + @classmethod + def setUpClass(cls): + if cls._imap_server_available(): + return + reason = ( + "No IMAP test server reachable at " + f"{cls.imap_host}:{cls.imap_port} - start the greenmail container " + "described in https://github.com/jan-janssen/testing-imap to run this test." + ) + if imap_integration_required(): + raise AssertionError( + f"{IMAP_INTEGRATION_REQUIRED} is set, so this test must actually run " + f"rather than be skipped. {reason}" + ) + raise unittest.SkipTest(reason) + + @classmethod + def _imap_server_available(cls, timeout=2.0, attempts=5, delay=1.5): + for attempt in range(attempts): + try: + with IMAP4(cls.imap_host, cls.imap_port, timeout=timeout) as client: + status, _ = client.noop() + if status == "OK": + return True + except OSError: + pass + if attempt < attempts - 1: + time.sleep(delay) + return False + + def setUp(self): + with IMAP4(self.imap_host, self.imap_port, timeout=10) as client: + client.login(self.username, self.password) + client.select("INBOX") + status, data = client.search(None, "ALL") + for message_id in data[0].split(): + client.store(message_id, "+FLAGS", r"(\Deleted)") + client.expunge() + for folder in ("MailSortInbox", "Sorted"): + client.create(folder) + + def _send_message(self, subject, body): + message_id = f"<{uuid.uuid4()}@example.test>" + message = EmailMessage() + message["From"] = "sender@example.test" + message["To"] = self.recipient + message["Subject"] = subject + message["Message-ID"] = message_id + message.set_content(body) + with smtplib.SMTP(self.smtp_host, self.smtp_port, timeout=10) as smtp: + smtp.send_message(message) + return message_id + + def _wait_for_message_in_inbox(self, message_id, timeout=10.0): + deadline = time.monotonic() + timeout + with IMAP4(self.imap_host, self.imap_port, timeout=10) as client: + client.login(self.username, self.password) + client.select("INBOX") + while time.monotonic() < deadline: + status, data = client.search( + None, "HEADER", "Message-ID", f'"{message_id}"' + ) + self.assertEqual(status, "OK") + if data[0].split(): + return + time.sleep(0.2) + self.fail(f"Message {message_id!r} was not delivered to INBOX") + + def test_update_database_and_move_round_trip(self): + message_id = self._send_message( + subject="Integration test message", + body="Body from gmailsorter IMAP test.", + ) + self._wait_for_message_in_inbox(message_id) + + with Imap( + host=self.imap_host, + port=self.imap_port, + username=self.username, + password=self.password, + connection_str="sqlite:///:memory:", + use_ssl=False, + ) as imap: + self._assert_move_round_trip(imap=imap) + + def _assert_move_round_trip(self, imap): + imap.update_database(quick=False) + df = imap.get_all_emails_in_database() + + self.assertIn("Integration test message", df["subject"].tolist()) + stored_id = df.loc[df["subject"] == "Integration test message", "id"].iloc[0] + self.assertTrue(stored_id.startswith("INBOX\x1f")) + + imap._modify_message_labels( + message_id=stored_id, + label_id_remove_lst=["INBOX"], + label_id_add_lst=["MailSortInbox"], + ) + + imap.update_database(quick=False) + df_after_move = imap.get_all_emails_in_database() + moved_row = df_after_move.loc[ + df_after_move["subject"] == "Integration test message" + ] + self.assertEqual(len(moved_row), 1) + self.assertTrue(moved_row.iloc[0]["id"].startswith("MailSortInbox\x1f")) + + +class TestImapIntegrationRequiredGate(unittest.TestCase): + """Covers the env-var gate itself, which needs no IMAP server.""" + + def test_not_required_by_default(self): + with patch.dict(os.environ, {}, clear=True): + self.assertFalse(imap_integration_required()) + + def test_not_required_for_empty_value(self): + with patch.dict(os.environ, {IMAP_INTEGRATION_REQUIRED: " "}): + self.assertFalse(imap_integration_required()) + + def test_required_for_truthy_value(self): + with patch.dict(os.environ, {IMAP_INTEGRATION_REQUIRED: "true"}): + self.assertTrue(imap_integration_required()) + + def test_set_up_class_skips_without_server_by_default(self): + with ( + patch.dict(os.environ, {}, clear=True), + patch.object( + TestImapServiceIntegration, "_imap_server_available", return_value=False + ), + self.assertRaises(unittest.SkipTest), + ): + TestImapServiceIntegration.setUpClass() + + def test_set_up_class_fails_without_server_when_required(self): + with ( + patch.dict(os.environ, {IMAP_INTEGRATION_REQUIRED: "true"}), + patch.object( + TestImapServiceIntegration, "_imap_server_available", return_value=False + ), + self.assertRaises(AssertionError), + ): + TestImapServiceIntegration.setUpClass() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mail_base.py b/tests/test_mail_base.py new file mode 100644 index 0000000..63cd0a9 --- /dev/null +++ b/tests/test_mail_base.py @@ -0,0 +1,134 @@ +from unittest import TestCase +from unittest.mock import MagicMock, patch + +import pandas as pd + +from gmailsorter.base.mail import AbstractMailBox + + +class _StubMailBox(AbstractMailBox): + """Minimal concrete AbstractMailBox used to test the shared loop in isolation.""" + + def __init__(self, label_dict_fixture=None, **kwargs): + self.label_dict_fixture = label_dict_fixture or { + "Inbox": "Inbox", + "Spam": "Spam", + } + self.search_result = [] + self.message_detail_dict = {} + self.modify_calls = [] + self.labels_for_email_dict = {} + super().__init__(mail_service=MagicMock(), **kwargs) + + def _search_email_on_server( + self, query_string="", label_lst=None, only_message_ids=False + ): + return self.search_result + + def _get_message_detail(self, message_id, email_format=None, metadata_headers=None): + return self.message_detail_dict.get(message_id) + + def _get_label_translate_dict(self): + return self.label_dict_fixture + + def _modify_message_labels( + self, message_id, label_id_remove_lst=None, label_id_add_lst=None + ): + self.modify_calls.append((message_id, label_id_remove_lst, label_id_add_lst)) + + def _get_labels_for_email(self, message_id): + return self.labels_for_email_dict.get(message_id, []) + + def _parse_message(self, message): + return message + + +class AbstractMailBoxTest(TestCase): + def test_labels_property(self): + mailbox = _StubMailBox() + self.assertEqual(sorted(mailbox.labels), ["Inbox", "Spam"]) + + def test_download_emails_for_label(self): + mailbox = _StubMailBox() + mailbox.search_result = ["id1", "id2"] + mailbox.message_detail_dict = { + "id1": { + "id": "id1", + "threads": "t1", + "labels": [], + "to": [], + "from": None, + "cc": [], + "subject": "s1", + "content": "c1", + "date": None, + }, + "id2": None, + } + + df = mailbox.download_emails_for_label(label="Inbox") + + self.assertEqual(df["id"].tolist(), ["id1"]) + + def test_move_emails_skips_matching_or_none_labels(self): + mailbox = _StubMailBox() + + mailbox._move_emails( + move_email_dict={"id1": None, "id2": "Inbox", "id3": "Spam"}, + label_to_ignore="Inbox", + ) + + self.assertEqual(mailbox.modify_calls, [("id3", ["Inbox"], ["Spam"])]) + + def test_update_database_marks_missing_as_deleted(self): + db_email = MagicMock() + db_email.get_labels_to_update.return_value = (["new"], [], ["deleted"]) + mailbox = _StubMailBox(database_email=db_email) + mailbox.search_result = ["new"] + mailbox.message_detail_dict = { + "new": { + "id": "new", + "threads": "t", + "labels": [], + "to": [], + "from": None, + "cc": [], + "subject": "s", + "content": "c", + "date": None, + } + } + + mailbox.update_database(quick=False) + + db_email.mark_emails_as_deleted.assert_called_once_with( + message_id_lst=["deleted"], user_id=1 + ) + db_email.store_dataframe.assert_called_once() + + @patch("gmailsorter.base.mail.fit_machine_learning_models") + @patch("gmailsorter.base.mail.encode_df_for_machine_learning") + def test_fit_machine_learning_model_to_database(self, encode_mock, fit_mock): + db_email = MagicMock() + db_email.get_all_emails.return_value = pd.DataFrame( + [ + { + "id": "x", + "from": "a@b.com", + "to": [], + "cc": [], + "labels": [], + "threads": "t", + } + ] + ) + db_ml = MagicMock() + mailbox = _StubMailBox(database_email=db_email, database_ml=db_ml) + features = pd.DataFrame([{"email_id": "x", "f1": 1}]) + labels = pd.DataFrame([{"labels_Inbox": 1}]) + encode_mock.return_value = (features, labels) + fit_mock.return_value = {"Inbox": MagicMock()} + + mailbox.fit_machine_learning_model_to_database(n_estimators=5, max_features=2) + + db_ml.store_models.assert_called_once() diff --git a/tests/test_main.py b/tests/test_main.py index 2e44c04..5c719dd 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -105,7 +105,9 @@ def test_label_flag_triggers_filter(self, load_secrets_mock, gmail_cls): email_download_format="metadata", ) gmail_instance.filter_messages_from_server.assert_called_once_with( - label="Inbox", recommendation_ratio=0.9 + label="Inbox", + recommendation_ratio=0.9, + label_prefix="labels_Label_", ) gmail_instance.update_database.assert_not_called() diff --git a/tests/test_message.py b/tests/test_message.py index 3e220ef..3d4ce85 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -1,6 +1,10 @@ from unittest import TestCase from datetime import datetime -from gmailsorter.base.message import email_date_converter, AbstractMessage +from gmailsorter.base.message import ( + email_date_converter, + AbstractMessage, + strip_html_tags, +) class MessageTest(TestCase): @@ -76,3 +80,9 @@ def test_email_date_converter(self): datetime.strptime("24-01-2022", "%d-%m-%Y"), ) self.assertEqual(email_date_converter(None), None) + + def test_strip_html_tags(self): + self.assertEqual( + strip_html_tags("

Hello World

"), + "Hello World", + )