From e932190618e9978a3facb39958d335f858aac2b4 Mon Sep 17 00:00:00 2001 From: Nathan Moinvaziri Date: Fri, 6 Feb 2026 10:26:44 -0800 Subject: [PATCH] Add imap_session module for centralized OAuth2 refresh (#21) * Add .claude to gitignore * Add imap_session module for centralized OAuth2 refresh Create new imap_session.py module that combines imap_common and imap_oauth2 for unified session management: - ensure_connection(): Refreshes OAuth2 token and ensures connection health - ensure_folder_session(): Ensures connection AND folder selection Updated backup, migrate, and restore scripts to use the new module, simplifying the proactive token refresh pattern in process_batch functions. * Move is_token_expired_error to imap_oauth2 module Relocates the OAuth2 token expiration detection function to the oauth2-specific module where it logically belongs. * Add retry logic for auth errors during email processing Extract single-UID processing into dedicated functions with retry capability for recoverable authentication failures. Broadens auth error detection beyond token expiration to include session loss and authentication failures. * Cache Microsoft tenant discovery results Avoids repeated OpenID Connect discovery requests for the same email domain during OAuth2 authentication flows. Normalizes domain to lowercase for consistent cache keys. * Add tests for is_token_expired_error and is_auth_error Covers token expiration, session invalidation, and authentication failure detection with case-insensitive matching verification. Ensures non-auth errors like connection and timeout issues are correctly excluded. * Add tests for imap_session module Covers ensure_connection and ensure_folder_session functions including OAuth2 token refresh, folder selection, and error handling scenarios. Verifies connection change detection triggers folder re-selection. * Add tests for auth error re-raising in get_message_ids_in_folder Verifies that authentication errors are propagated to callers for reconnection handling while non-auth errors return empty results. Clarifies test names to distinguish error handling behaviors. --- .gitignore | 1 + src/backup_imap_emails.py | 233 +++++++++++++++------- src/imap_common.py | 12 +- src/imap_oauth2.py | 41 +++- src/imap_session.py | 62 ++++++ src/migrate_imap_emails.py | 337 +++++++++++++++++++------------- src/restore_imap_emails.py | 24 +-- test/test_backup_imap_emails.py | 16 +- test/test_imap_common.py | 29 ++- test/test_imap_oauth2.py | 100 ++++++++++ test/test_imap_session.py | 268 +++++++++++++++++++++++++ 11 files changed, 887 insertions(+), 236 deletions(-) create mode 100644 src/imap_session.py create mode 100644 test/test_imap_session.py diff --git a/.gitignore b/.gitignore index c037455..e760a40 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ env/ venv/ .env .coverage +.claude diff --git a/src/backup_imap_emails.py b/src/backup_imap_emails.py index 78409fc..ea53621 100644 --- a/src/backup_imap_emails.py +++ b/src/backup_imap_emails.py @@ -58,6 +58,7 @@ import threading import imap_common import imap_oauth2 +import imap_session # Defaults MAX_WORKERS = 10 @@ -76,13 +77,6 @@ def safe_print(message): print(f"[{short_name}] {message}") -def ensure_connection(conn, conf): - """Refresh OAuth2 token if needed and ensure connection is healthy.""" - if conf.get("oauth2"): - imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token")) - return imap_common.ensure_connection_from_conf(conn, conf) - - def get_thread_connection(src_conf): # Backwards-compatible: tests and some internal call sites may still pass # (host, user, password) tuples. Normalize to the dict form used by @@ -90,10 +84,74 @@ def get_thread_connection(src_conf): if isinstance(src_conf, tuple) and len(src_conf) == 3: src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]} - thread_local.src = ensure_connection(getattr(thread_local, "src", None), src_conf) + thread_local.src = imap_session.ensure_connection(getattr(thread_local, "src", None), src_conf) return thread_local.src +def process_single_uid(src, uid, folder_name, local_folder_path): + """ + Fetch and save a single email by UID. + + Args: + src: IMAP connection + uid: Message UID to fetch + folder_name: Current folder name (for logging) + local_folder_path: Local directory to save email + + Returns: + Tuple of (success: bool, connection): + - (True, src): UID processed successfully (or skipped) + - (False, src): Auth error occurred, caller should retry after reconnect + """ + uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid) + + try: + resp, data = src.uid("fetch", uid, "(RFC822)") + if resp != "OK" or not data or data[0] is None: + safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}") + return (True, src) # Don't retry, move on + + raw_email = None + for item in data: + if isinstance(item, tuple): + raw_email = item[1] + break + + # Derive Subject for filename from the already-fetched message bytes. + _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email) + if not subject: + clean_subject = "No Subject" + else: + clean_subject = imap_common.sanitize_filename(subject) + clean_subject = clean_subject[:100] + + filename = f"{uid_str}_{clean_subject}.eml" + full_path = os.path.join(local_folder_path, filename) + + if os.path.exists(full_path): + return (True, src) # Already exists + + if raw_email: + try: + with open(full_path, "wb") as f: + f.write(raw_email) + safe_print(f"[{folder_name}] SAVED | {filename[:60]}...") + except OSError as e: + safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}") + else: + safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}") + + return (True, src) + + except Exception as e: + if imap_oauth2.is_auth_error(e): + safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...") + return (False, src) # Signal retry needed + else: + safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}") + return (True, src) # Don't retry other errors + + def process_batch(uids, folder_name, src_conf, local_folder_path): if isinstance(src_conf, tuple) and len(src_conf) == 3: src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]} @@ -110,52 +168,24 @@ def process_batch(uids, folder_name, src_conf, local_folder_path): return for uid in uids: - try: - # Helper to handle byte UIDs - uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid) + uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid) + max_retries = 2 - # Fetch Full Content - # RFC822 gets the whole message including headers and attachments - resp, data = src.uid("fetch", uid, "(RFC822)") - if resp != "OK" or not data or data[0] is None: - safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}") - continue + for attempt in range(max_retries): + src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True) + thread_local.src = src + if not ok: + safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}") + return - raw_email = None - for item in data: - if isinstance(item, tuple): - raw_email = item[1] - break + success, src = process_single_uid(src, uid, folder_name, local_folder_path) + thread_local.src = src - # Derive Subject for filename from the already-fetched message bytes. - _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email) - if not subject: - clean_subject = "No Subject" - else: - clean_subject = imap_common.sanitize_filename(subject) - # Limit subject len to avoid FS path length limits (max 255 usually) - clean_subject = clean_subject[:100] - - filename = f"{uid_str}_{clean_subject}.eml" - full_path = os.path.join(local_folder_path, filename) - - # Double check existence (in case incremental UID checks missed it or naming collision). - if os.path.exists(full_path): - continue - - if raw_email: - try: - with open(full_path, "wb") as f: - f.write(raw_email) - safe_print(f"[{folder_name}] SAVED | {filename[:60]}...") - except OSError as e: - # Valid filename might still fail if path too long - safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}") - else: - safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}") - - except Exception as e: - safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}") + if success: + break + if attempt < max_retries - 1: + src = None + thread_local.src = None def get_existing_uids(local_path): @@ -210,11 +240,21 @@ def is_gmail_label_folder(folder_name): PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS -def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None): +def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None): """ - Returns a dict of Message-IDs and their IMAP flags for all emails in a folder. - Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... } - Optional progress_callback(current, total) for progress reporting. + Returns a dict of Message-IDs and their IMAP flags for all emails in a folder, + with OAuth2 session management. + + Args: + imap_conn: IMAP connection + folder_name: Folder to scan + src_conf: Connection config dict for OAuth2 refresh + progress_callback: Optional callback(current, total) for progress reporting + + Returns: + Tuple of (message_info, imap_conn) where message_info is: + { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... } + The returned imap_conn may be different if reconnection occurred. """ message_info = {} @@ -222,16 +262,16 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None): imap_conn.select(f'"{folder_name}"', readonly=True) except Exception as e: safe_print(f"Could not select folder {folder_name}: {e}") - return message_info + return (message_info, imap_conn) try: resp, data = imap_conn.uid("search", None, "ALL") if resp != "OK" or not data or not data[0]: - return message_info + return (message_info, imap_conn) uids = data[0].split() if not uids: - return message_info + return (message_info, imap_conn) total_uids = len(uids) @@ -245,6 +285,13 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None): if progress_callback: progress_callback(min(i + batch_size, total_uids), total_uids) + # Proactively refresh token and ensure folder is selected + if src_conf: + imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True) + if not ok: + safe_print(f"ERROR: Connection/folder lost in {folder_name}") + break + try: # Fetch both Message-ID header and FLAGS resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])") @@ -271,19 +318,44 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None): message_info[msg_id] = {"flags": flags} except Exception as e: safe_print(f"Error fetching batch in {folder_name}: {e}") - # Try to keep connection alive - try: - imap_conn.noop() - except Exception: - pass continue except Exception as e: safe_print(f"Error searching folder {folder_name}: {e}") + return (message_info, imap_conn) + + +def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None): + """ + Returns a dict of Message-IDs and their IMAP flags for all emails in a folder. + + Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... } + Optional progress_callback(current, total) for progress reporting. + """ + message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback) return message_info +def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None): + """ + Returns a set of Message-IDs for all emails in a given folder, + with OAuth2 session management. + + Args: + imap_conn: IMAP connection + folder_name: Folder to scan + src_conf: Connection config dict for OAuth2 refresh + progress_callback: Optional callback(current, total) for progress reporting + + Returns: + Tuple of (message_ids, imap_conn) where message_ids is a set of strings. + The returned imap_conn may be different if reconnection occurred. + """ + info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback) + return (set(info.keys()), imap_conn) + + def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None): """ Returns a set of Message-IDs for all emails in a given folder. @@ -294,7 +366,7 @@ def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None): return set(info.keys()) -def build_labels_manifest(imap_conn, local_path): +def build_labels_manifest(imap_conn, local_path, src_conf=None): """ Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags. Scans all folders (labels) in the account and records which Message-IDs @@ -302,6 +374,7 @@ def build_labels_manifest(imap_conn, local_path): Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... } Saves the manifest to labels_manifest.json in the backup directory. + Optional src_conf for automatic token refresh on expiration. """ import time @@ -329,7 +402,9 @@ def build_labels_manifest(imap_conn, local_path): flush=True, ) - all_mail_info = get_message_info_in_folder(imap_conn, imap_common.GMAIL_ALL_MAIL, all_mail_progress_cb) + all_mail_info, imap_conn = get_message_info_in_folder_with_conf( + imap_conn, imap_common.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb + ) print() # New line after progress # Initialize manifest with flags from All Mail @@ -369,7 +444,7 @@ def build_labels_manifest(imap_conn, local_path): ) safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}") - message_ids = get_message_ids_in_folder(imap_conn, folder_name, progress_cb) + message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb) print() # New line after progress # Keep connection alive between folders @@ -422,13 +497,14 @@ def build_labels_manifest(imap_conn, local_path): return manifest -def build_flags_manifest(imap_conn, local_path, folders_to_scan=None): +def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None): """ Builds a manifest mapping Message-IDs to their IMAP flags. For non-Gmail servers, scans specified folders (or all folders if not specified). Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... } Saves the manifest to flags_manifest.json in the backup directory. + Optional src_conf for automatic token refresh on expiration. """ import time @@ -467,7 +543,7 @@ def build_flags_manifest(imap_conn, local_path, folders_to_scan=None): ) safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}") - folder_info = get_message_info_in_folder(imap_conn, folder_name, progress_cb) + folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb) print() # New line after progress # Merge into manifest (keep first occurrence of flags) @@ -829,7 +905,7 @@ def main(): if args.preserve_labels or args.manifest_only or args.gmail_mode: print("Building Gmail labels manifest...") print("This scans all folders to map Message-IDs to labels and flags.\n") - build_labels_manifest(src, local_path) + build_labels_manifest(src, local_path, src_conf) print("") # Blank line after manifest building # Build flags-only manifest for non-Gmail servers elif args.preserve_flags: @@ -837,12 +913,15 @@ def main(): print("This scans folders to capture read/starred/etc status.\n") # Get folders to scan folders_to_scan = [args.folder] if args.folder else None - build_flags_manifest(src, local_path, folders_to_scan) + build_flags_manifest(src, local_path, folders_to_scan, src_conf) print("") # Blank line after manifest building # If manifest-only mode, we're done if args.manifest_only: - src.logout() + try: + src.logout() + except Exception: + pass # Connection may already be closed manifest_path = os.path.join(local_path, MANIFEST_FILENAME) print("\nManifest-only mode complete.") print(f"Labels manifest saved to: {manifest_path}") @@ -856,15 +935,23 @@ def main(): elif args.folder: backup_folder(src, args.folder, local_path, src_conf, args.dest_delete) else: + # Reconnect after potentially long manifest building + src = imap_session.ensure_connection(src, src_conf) + if not src: + print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.") + sys.exit(0) folders = imap_common.list_selectable_folders(src) for name in folders: - src = ensure_connection(src, src_conf) + src = imap_session.ensure_connection(src, src_conf) if not src: print("Fatal: Could not reconnect to IMAP server. Aborting.") sys.exit(1) backup_folder(src, name, local_path, src_conf, args.dest_delete) - src.logout() + try: + src.logout() + except Exception: + pass # Connection may already be closed print("\nBackup completed successfully.") if args.preserve_labels or args.gmail_mode: diff --git a/src/imap_common.py b/src/imap_common.py index 2e146bd..6fe7434 100644 --- a/src/imap_common.py +++ b/src/imap_common.py @@ -14,6 +14,8 @@ from email import policy from email.header import decode_header from email.parser import BytesParser +import imap_oauth2 + # Standard IMAP flags FLAG_SEEN = "\\Seen" FLAG_ANSWERED = "\\Answered" @@ -363,7 +365,10 @@ def get_message_ids_in_folder(imap_conn): resp, data = imap_conn.uid("search", None, "ALL") if resp != "OK" or not data[0].strip(): return {} - except Exception: + except Exception as e: + # Re-raise token expiration errors so callers can handle reconnection + if imap_oauth2.is_auth_error(e): + raise return {} uids = data[0].split() @@ -403,7 +408,10 @@ def get_uid_to_message_id_map(imap_conn, uids): mid = extract_message_id(item[1]) if mid: uid_to_msgid[uid] = mid - except Exception: + except Exception as e: + # Re-raise token expiration errors so callers can handle reconnection + if imap_oauth2.is_auth_error(e): + raise continue return uid_to_msgid diff --git a/src/imap_oauth2.py b/src/imap_oauth2.py index a366f1c..870e030 100644 --- a/src/imap_oauth2.py +++ b/src/imap_oauth2.py @@ -23,6 +23,7 @@ import urllib.parse # Module-level caches for OAuth2 token refresh _msal_app_cache = {} # (client_id, tenant_id) -> PublicClientApplication _google_creds_cache = {} # (client_id, client_secret) -> credentials +_tenant_cache = {} # domain -> tenant_id _token_refresh_lock = threading.Lock() @@ -46,6 +47,35 @@ def _fetch_json_https(host, path, timeout=10): return json.loads(body.decode("utf-8")) +def is_token_expired_error(error): + """ + Check if an exception indicates OAuth2 token expiration. + + Args: + error: The exception to check + + Returns: + True if the error indicates token expiration, False otherwise + """ + error_str = str(error).lower() + return "accesstokenexpired" in error_str or "session invalidated" in error_str + + +def is_auth_error(error): + """ + Check if an exception indicates an authentication failure that may be recoverable + by reconnecting (token expiration, session loss, etc.). + + Args: + error: The exception to check + + Returns: + True if the error indicates an auth failure, False otherwise + """ + error_str = str(error).lower() + return is_token_expired_error(error) or "not authenticated" in error_str or "authentication failed" in error_str + + def detect_oauth2_provider(host): """ Detects the OAuth2 provider from the IMAP host. @@ -63,13 +93,18 @@ def discover_microsoft_tenant(email): """ Auto-discovers the Microsoft tenant ID from an email address domain. Uses the OpenID Connect discovery endpoint (no authentication required). + Results are cached per domain to avoid repeated network requests. Returns the tenant ID string or None if discovery fails. """ - domain = email.split("@")[-1].strip() + domain = email.split("@")[-1].strip().lower() if not domain: print("Error: Could not discover Microsoft tenant: missing email domain") return None + # Return cached tenant if available + if domain in _tenant_cache: + return _tenant_cache[domain] + domain_quoted = urllib.parse.quote(domain, safe=".-") path = f"/{domain_quoted}/.well-known/openid-configuration" @@ -82,7 +117,9 @@ def discover_microsoft_tenant(email): issuer = data.get("issuer", "") match = re.search(r"/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})", issuer) if match: - return match.group(1) + tenant_id = match.group(1) + _tenant_cache[domain] = tenant_id + return tenant_id print(f"Error: Could not extract tenant ID from issuer: {issuer}") return None diff --git a/src/imap_session.py b/src/imap_session.py new file mode 100644 index 0000000..c86193b --- /dev/null +++ b/src/imap_session.py @@ -0,0 +1,62 @@ +""" +IMAP Session Management + +Connection and session management for IMAP operations with OAuth2 support. +Combines imap_common (low-level IMAP) with imap_oauth2 (token refresh). +""" + +import imap_common +import imap_oauth2 + + +def ensure_connection(conn, conf): + """ + Refresh OAuth2 token if needed and ensure connection is healthy. + + Args: + conn: Existing IMAP connection or None + conf: Connection config dict with keys: + - host, user, password, oauth2_token: connection credentials + - oauth2: dict with provider, client_id, email, client_secret (optional) + + Returns: + Healthy IMAP connection, or None if connection failed. + May return a different connection object if reconnection was needed. + """ + if conf.get("oauth2"): + imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token")) + return imap_common.ensure_connection_from_conf(conn, conf) + + +def ensure_folder_session(conn, conf, folder_name, readonly=True): + """ + Ensure connection is healthy and folder is selected. + + Proactively refreshes OAuth2 token if needed. If the connection was + refreshed (new connection object), reselects the folder. + + Args: + conn: Existing IMAP connection or None + conf: Connection config dict (see ensure_connection) + folder_name: IMAP folder to select + readonly: Whether to select folder as readonly (default True) + + Returns: + Tuple of (connection, success: bool) + - On success: (connection, True) - folder is selected + - On failure: (connection or None, False) + """ + old_conn = conn + new_conn = ensure_connection(conn, conf) + + if not new_conn: + return None, False + + # If connection changed (token refreshed), need to reselect folder + if new_conn is not old_conn: + try: + new_conn.select(f'"{folder_name}"', readonly=readonly) + except Exception: + return new_conn, False + + return new_conn, True diff --git a/src/migrate_imap_emails.py b/src/migrate_imap_emails.py index 398c9ca..b55dbab 100644 --- a/src/migrate_imap_emails.py +++ b/src/migrate_imap_emails.py @@ -117,6 +117,7 @@ import threading import imap_common import imap_oauth2 +import imap_session # Configuration defaults DELETE_FROM_SOURCE_DEFAULT = False @@ -288,19 +289,167 @@ def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msg return deleted_count -def ensure_connection(conn, conf): - """Refresh OAuth2 token if needed and ensure connection is healthy.""" - if conf.get("oauth2"): - imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token")) - return imap_common.ensure_connection_from_conf(conn, conf) - - def get_thread_connections(src_conf, dest_conf): - thread_local.src = ensure_connection(getattr(thread_local, "src", None), src_conf) - thread_local.dest = ensure_connection(getattr(thread_local, "dest", None), dest_conf) + thread_local.src = imap_session.ensure_connection(getattr(thread_local, "src", None), src_conf) + thread_local.dest = imap_session.ensure_connection(getattr(thread_local, "dest", None), dest_conf) return thread_local.src, thread_local.dest +def process_single_uid( + src, + dest, + uid, + folder_name, + delete_from_source, + trash_folder, + preserve_flags, + gmail_mode, + label_index, + check_duplicate, +): + """ + Migrate a single email by UID. + + Returns: + Tuple of (success, src, dest, deleted): + - success=True: UID processed (copied, skipped, or non-auth error) + - success=False: Auth error, caller should retry after reconnect + - deleted: 1 if message was marked for deletion, 0 otherwise + """ + uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid) + + try: + resp, data = src.uid("fetch", uid, "(FLAGS INTERNALDATE BODY.PEEK[])") + if resp != "OK": + safe_print(f"[{folder_name}] ERROR Fetch | UID {uid_str}") + return (True, src, dest, 0) + + msg_content = None + flags = None + date_str = None + + for item in data: + if isinstance(item, tuple): + msg_content = item[1] + meta = item[0].decode("utf-8", errors="ignore") + flags_match = re.search(r"FLAGS\s+\((.*?)\)", meta) + if flags_match: + flags = filter_preservable_flags(flags_match.group(1)) + date_match = re.search(r"INTERNALDATE\s+\"(.*?)\"", meta) + if date_match: + date_str = f'"{date_match.group(1)}"' + + if not msg_content: + return (True, src, dest, 0) + + size = len(msg_content) if isinstance(msg_content, (bytes, bytearray)) else 0 + size_str = f"{size / 1024:.1f}KB" if size else "0KB" + msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(msg_content) + + if not msg_id: + msg_id = imap_common.parse_message_id_from_bytes(msg_content) + + # Determine target folder and labels for Gmail mode + apply_labels = gmail_mode + labels = [] + if apply_labels and msg_id and label_index is not None: + labels = sorted(label_index.get(msg_id, set())) + + if apply_labels: + skip_folders = {imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH} + target_folder = None + remaining_labels = [] + + for label in labels: + label_folder = label_to_folder(label) + if label_folder in skip_folders: + continue + if target_folder is None: + target_folder = label_folder + else: + remaining_labels.append(label) + + if target_folder is None: + target_folder = imap_common.FOLDER_RESTORED_UNLABELED + remaining_labels = [] + else: + target_folder = folder_name + + imap_common.ensure_folder_exists(dest, target_folder) + dest.select(f'"{target_folder}"') + + is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id)) + + if is_duplicate: + safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}") + if preserve_flags and flags and msg_id: + sync_flags_on_existing(dest, target_folder, msg_id, flags, size) + else: + valid_flags = f"({flags})" if (preserve_flags and flags) else None + imap_common.append_email( + dest, + target_folder, + msg_content, + date_str, + valid_flags, + ensure_folder=False, + ) + safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}") + if preserve_flags and flags: + for flag in flags.split(): + safe_print(f" -> Applied flag: {flag}") + + # Apply remaining Gmail labels + if apply_labels and remaining_labels and msg_id: + for label in remaining_labels: + label_folder = label_to_folder(label) + if label_folder == target_folder: + continue + if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH): + continue + try: + imap_common.ensure_folder_exists(dest, label_folder) + dest.select(f'"{label_folder}"') + if not imap_common.message_exists_in_folder(dest, msg_id): + valid_flags = f"({flags})" if (preserve_flags and flags) else None + imap_common.append_email( + dest, + label_folder, + msg_content, + date_str, + valid_flags, + ensure_folder=False, + ) + safe_print(f" -> Applied label: {label}") + if preserve_flags and flags: + for flag in flags.split(): + safe_print(f" -> Applied flag: {flag}") + elif preserve_flags and flags: + sync_flags_on_existing(dest, label_folder, msg_id, flags, size) + except Exception as e: + safe_print(f" -> Error applying label {label}: {e}") + + deleted = 0 + if delete_from_source: + if trash_folder and folder_name != trash_folder: + try: + src.uid("copy", uid, f'"{trash_folder}"') + except Exception: + pass + src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL) + deleted = 1 + + return (True, src, dest, deleted) + + except Exception as e: + if imap_oauth2.is_auth_error(e): + safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...") + return (False, src, dest, 0) + else: + safe_print(f"[{folder_name}] ERROR Exec | UID {uid_str}: {e}") + return (True, src, dest, 0) + + def process_batch( uids, folder_name, @@ -318,14 +467,12 @@ def process_batch( safe_print("Error: Could not establish connections in worker thread.") return - # Select source folder try: src.select(f'"{folder_name}"', readonly=False) except Exception as e: safe_print(f"Error selecting folder {folder_name} in worker: {e}") return - # In non-Gmail-mode, we keep a selected destination folder for efficiency if not gmail_mode: try: imap_common.ensure_folder_exists(dest, folder_name) @@ -335,134 +482,54 @@ def process_batch( return deleted_count = 0 + for uid in uids: - try: - # Fetch full message (needed to copy and/or apply labels) - resp, data = src.uid("fetch", uid, "(FLAGS INTERNALDATE BODY.PEEK[])") - if resp != "OK": - uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid) - safe_print(f"[{folder_name}] ERROR Fetch | UID {uid_str}") - continue + uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid) + max_retries = 2 - msg_content = None - flags = None - date_str = None + for attempt in range(max_retries): + src, src_ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=False) + thread_local.src = src + if not src_ok: + safe_print(f"[{folder_name}] ERROR: Source connection/folder lost for UID {uid_str}") + return - for item in data: - if isinstance(item, tuple): - msg_content = item[1] - meta = item[0].decode("utf-8", errors="ignore") - flags_match = re.search(r"FLAGS\s+\((.*?)\)", meta) - if flags_match: - flags = filter_preservable_flags(flags_match.group(1)) - date_match = re.search(r"INTERNALDATE\s+\"(.*?)\"", meta) - if date_match: - date_str = f'"{date_match.group(1)}"' - - if not msg_content: - continue - - # Compute size and parse headers from the already-fetched message bytes. - size = len(msg_content) if isinstance(msg_content, (bytes, bytearray)) else 0 - size_str = f"{size / 1024:.1f}KB" if size else "0KB" - msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(msg_content) - - # Prefer Message-ID from body if missing from header parse - if not msg_id: - msg_id = imap_common.parse_message_id_from_bytes(msg_content) - - # Determine target folder and labels for Gmail mode - apply_labels = gmail_mode - labels = [] - if apply_labels and msg_id and label_index is not None: - labels = sorted(label_index.get(msg_id, set())) - - if apply_labels: - skip_folders = {imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH} - target_folder = None - remaining_labels = [] - - for label in labels: - label_folder = label_to_folder(label) - if label_folder in skip_folders: - continue - if target_folder is None: - target_folder = label_folder - else: - remaining_labels.append(label) - - if target_folder is None: - target_folder = imap_common.FOLDER_RESTORED_UNLABELED - remaining_labels = [] + if not gmail_mode: + dest, dest_ok = imap_session.ensure_folder_session(dest, dest_conf, folder_name, readonly=False) + thread_local.dest = dest + if not dest_ok: + safe_print(f"[{folder_name}] ERROR: Dest connection/folder lost for UID {uid_str}") + return else: - target_folder = folder_name - remaining_labels = [] + dest = imap_session.ensure_connection(dest, dest_conf) + thread_local.dest = dest + if not dest: + safe_print(f"[{folder_name}] ERROR: Dest connection lost for UID {uid_str}") + return - imap_common.ensure_folder_exists(dest, target_folder) - dest.select(f'"{target_folder}"') + success, src, dest, deleted = process_single_uid( + src, + dest, + uid, + folder_name, + delete_from_source, + trash_folder, + preserve_flags, + gmail_mode, + label_index, + check_duplicate, + ) + thread_local.src = src + thread_local.dest = dest + deleted_count += deleted - is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id)) - - if is_duplicate: - safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}") - if preserve_flags and flags and msg_id: - sync_flags_on_existing(dest, target_folder, msg_id, flags, size) - else: - valid_flags = f"({flags})" if (preserve_flags and flags) else None - imap_common.append_email( - dest, - target_folder, - msg_content, - date_str, - valid_flags, - ensure_folder=False, - ) - safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}") - if preserve_flags and flags: - for flag in flags.split(): - safe_print(f" -> Applied flag: {flag}") - - # Apply remaining Gmail labels (always, whether copied or skipped) - if apply_labels and remaining_labels and msg_id: - for label in remaining_labels: - label_folder = label_to_folder(label) - if label_folder == target_folder: - continue - if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH): - continue - try: - imap_common.ensure_folder_exists(dest, label_folder) - dest.select(f'"{label_folder}"') - if not imap_common.message_exists_in_folder(dest, msg_id): - valid_flags = f"({flags})" if (preserve_flags and flags) else None - imap_common.append_email( - dest, - label_folder, - msg_content, - date_str, - valid_flags, - ensure_folder=False, - ) - safe_print(f" -> Applied label: {label}") - if preserve_flags and flags: - for flag in flags.split(): - safe_print(f" -> Applied flag: {flag}") - elif preserve_flags and flags: - sync_flags_on_existing(dest, label_folder, msg_id, flags, size) - except Exception as e: - safe_print(f" -> Error applying label {label}: {e}") - - if delete_from_source: - if trash_folder and folder_name != trash_folder: - try: - src.uid("copy", uid, f'"{trash_folder}"') - except Exception: - pass - src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL) - deleted_count += 1 - - except Exception as e: - safe_print(f"[{folder_name}] ERROR Exec | UID {uid}: {e}") + if success: + break + if attempt < max_retries - 1: + src = None + dest = None + thread_local.src = None + thread_local.dest = None if delete_from_source and deleted_count > 0: try: @@ -961,11 +1028,11 @@ def main(): safe_print(f"Skipping migration of Trash folder '{name}' (preventing circular migration).") continue - src_main = ensure_connection(src_main, src_conf) + src_main = imap_session.ensure_connection(src_main, src_conf) if not src_main: safe_print("Fatal: Could not reconnect to source IMAP server. Aborting.") sys.exit(1) - dest_main = ensure_connection(dest_main, dest_conf) + dest_main = imap_session.ensure_connection(dest_main, dest_conf) if not dest_main: safe_print("Fatal: Could not reconnect to destination IMAP server. Aborting.") sys.exit(1) diff --git a/src/restore_imap_emails.py b/src/restore_imap_emails.py index 8c3e826..b8c3476 100644 --- a/src/restore_imap_emails.py +++ b/src/restore_imap_emails.py @@ -61,6 +61,7 @@ from typing import Optional import imap_common import imap_oauth2 +import imap_session import restore_cache @@ -88,16 +89,9 @@ def safe_print(message): print(f"[{short_name}] {message}") -def ensure_connection(conn, conf): - """Refresh OAuth2 token if needed and ensure connection is healthy.""" - if conf.get("oauth2"): - imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token")) - return imap_common.ensure_connection_from_conf(conn, conf) - - def get_thread_connection(dest_conf): """Get or create a thread-local IMAP connection.""" - thread_local.dest = ensure_connection(getattr(thread_local, "dest", None), dest_conf) + thread_local.dest = imap_session.ensure_connection(getattr(thread_local, "dest", None), dest_conf) return thread_local.dest @@ -395,10 +389,16 @@ def process_restore_batch( gmail_mode = folder_name == "__GMAIL_MODE__" for file_path, filename in eml_files: + # Proactively refresh token if needed + dest = get_thread_connection(dest_conf) + if not dest: + safe_print(f"ERROR: Connection lost for {filename}") + return + try: message_id, date_str, raw_content, subject = parse_eml_file(file_path) if raw_content is None: - continue + continue # No content, skip to next file size = len(raw_content) size_str = f"{size / 1024:.1f}KB" @@ -472,7 +472,7 @@ def process_restore_batch( and message_id in existing_dest_msg_ids ): safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}") - continue + continue # Skip to next file # Upload to target folder. # Keep server-side duplicate checks enabled to avoid creating duplicates for emails @@ -815,7 +815,7 @@ def restore_folder( # Delete orphan emails from destination if enabled if dest_delete and local_msg_ids is not None: safe_print("Syncing destination: removing emails not in local backup...") - dest = ensure_connection(None, dest_conf) + dest = imap_session.ensure_connection(None, dest_conf) if dest: delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids) dest.logout() @@ -1192,7 +1192,7 @@ def main(): continue # Proactively refresh OAuth2 token and ensure connection is healthy between folders - dest = ensure_connection(dest, dest_conf) + dest = imap_session.ensure_connection(dest, dest_conf) if not dest: print("Fatal: Could not reconnect to destination IMAP server. Aborting.") sys.exit(1) diff --git a/test/test_backup_imap_emails.py b/test/test_backup_imap_emails.py index 3d88a10..a758dcf 100644 --- a/test/test_backup_imap_emails.py +++ b/test/test_backup_imap_emails.py @@ -551,16 +551,16 @@ class TestGmailLabelsPreservation: "": {"flags": []}, } - def mock_get_message_ids(conn, folder, progress_cb=None): - return folder_data.get(folder, set()) - - def mock_get_message_info(conn, folder, progress_cb=None): + def mock_get_message_info_with_conf(conn, folder, src_conf, progress_cb=None): if folder == "[Gmail]/All Mail": - return all_mail_info - return {} + return (all_mail_info, conn) + return ({}, conn) - monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder", mock_get_message_ids) - monkeypatch.setattr(backup_imap_emails, "get_message_info_in_folder", mock_get_message_info) + def mock_get_message_ids_with_conf(conn, folder, src_conf, progress_cb=None): + return (folder_data.get(folder, set()), conn) + + monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder_with_conf", mock_get_message_ids_with_conf) + monkeypatch.setattr(backup_imap_emails, "get_message_info_in_folder_with_conf", mock_get_message_info_with_conf) result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path)) diff --git a/test/test_imap_common.py b/test/test_imap_common.py index fca563e..0a0d9f6 100644 --- a/test/test_imap_common.py +++ b/test/test_imap_common.py @@ -16,6 +16,8 @@ import os import sys from unittest.mock import Mock +import pytest + sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) import imap_common @@ -653,14 +655,22 @@ class TestGetMessageIdsInFolder: result = imap_common.get_message_ids_in_folder(mock_conn) assert result == {} - def test_search_exception(self): - """Test returns empty dict when search raises exception.""" + def test_search_non_auth_exception_returns_empty(self): + """Test returns empty dict when search raises non-auth exception.""" mock_conn = Mock() mock_conn.uid.side_effect = Exception("Connection error") result = imap_common.get_message_ids_in_folder(mock_conn) assert result == {} + def test_search_auth_error_reraises(self): + """Test re-raises auth errors so callers can handle reconnection.""" + mock_conn = Mock() + mock_conn.uid.side_effect = Exception("User not authenticated") + + with pytest.raises(Exception, match="not authenticated"): + imap_common.get_message_ids_in_folder(mock_conn) + def test_single_message(self): """Test fetching single message ID.""" mock_conn = Mock() @@ -707,8 +717,8 @@ class TestGetMessageIdsInFolder: result = imap_common.get_message_ids_in_folder(mock_conn) assert result == {} - def test_fetch_exception_for_batch(self): - """Test continues when fetch raises exception for a batch.""" + def test_fetch_non_auth_exception_continues(self): + """Test continues when fetch raises non-auth exception for a batch.""" mock_conn = Mock() mock_conn.uid.side_effect = [ ("OK", [b"1"]), @@ -718,6 +728,17 @@ class TestGetMessageIdsInFolder: result = imap_common.get_message_ids_in_folder(mock_conn) assert result == {} + def test_fetch_auth_error_reraises(self): + """Test re-raises auth errors during fetch so callers can handle reconnection.""" + mock_conn = Mock() + mock_conn.uid.side_effect = [ + ("OK", [b"1"]), + Exception("AccessTokenExpired"), + ] + + with pytest.raises(Exception, match="AccessTokenExpired"): + imap_common.get_message_ids_in_folder(mock_conn) + def test_skips_empty_message_id(self): """Test that empty message IDs are not added to dict.""" mock_conn = Mock() diff --git a/test/test_imap_oauth2.py b/test/test_imap_oauth2.py index fe8e766..6a8e48d 100644 --- a/test/test_imap_oauth2.py +++ b/test/test_imap_oauth2.py @@ -28,9 +28,11 @@ def clear_oauth2_caches(): """Clear module-level OAuth2 caches between tests.""" imap_oauth2._msal_app_cache.clear() imap_oauth2._google_creds_cache.clear() + imap_oauth2._tenant_cache.clear() yield imap_oauth2._msal_app_cache.clear() imap_oauth2._google_creds_cache.clear() + imap_oauth2._tenant_cache.clear() class TestDetectOauth2Provider: @@ -487,3 +489,101 @@ class TestRefreshOauth2Token: # The other two should see conf["oauth2_token"] changed and skip. assert call_count["value"] == 1 assert conf["oauth2_token"] == "fresh_token" + + +class TestIsTokenExpiredError: + """Tests for is_token_expired_error function.""" + + def test_detects_access_token_expired(self): + """Test detects 'accesstokenexpired' error.""" + error = Exception("AUTHENTICATE failed: AccessTokenExpired") + assert imap_oauth2.is_token_expired_error(error) is True + + def test_detects_session_invalidated(self): + """Test detects 'session invalidated' error.""" + error = Exception("Session invalidated by server") + assert imap_oauth2.is_token_expired_error(error) is True + + def test_case_insensitive_access_token_expired(self): + """Test case-insensitive matching for accesstokenexpired.""" + error = Exception("ACCESSTOKENEXPIRED") + assert imap_oauth2.is_token_expired_error(error) is True + + def test_case_insensitive_session_invalidated(self): + """Test case-insensitive matching for session invalidated.""" + error = Exception("SESSION INVALIDATED") + assert imap_oauth2.is_token_expired_error(error) is True + + def test_false_for_connection_error(self): + """Test returns False for connection errors.""" + error = Exception("Connection refused") + assert imap_oauth2.is_token_expired_error(error) is False + + def test_false_for_timeout_error(self): + """Test returns False for timeout errors.""" + error = Exception("The read operation timed out") + assert imap_oauth2.is_token_expired_error(error) is False + + def test_false_for_generic_error(self): + """Test returns False for generic errors.""" + error = Exception("Something went wrong") + assert imap_oauth2.is_token_expired_error(error) is False + + def test_false_for_not_authenticated(self): + """Test returns False for 'not authenticated' (handled by is_auth_error).""" + error = Exception("User not authenticated") + assert imap_oauth2.is_token_expired_error(error) is False + + +class TestIsAuthError: + """Tests for is_auth_error function.""" + + def test_detects_access_token_expired(self): + """Test detects 'accesstokenexpired' error (via is_token_expired_error).""" + error = Exception("AUTHENTICATE failed: AccessTokenExpired") + assert imap_oauth2.is_auth_error(error) is True + + def test_detects_session_invalidated(self): + """Test detects 'session invalidated' error (via is_token_expired_error).""" + error = Exception("Session invalidated by server") + assert imap_oauth2.is_auth_error(error) is True + + def test_detects_not_authenticated(self): + """Test detects 'not authenticated' error.""" + error = Exception("User not authenticated") + assert imap_oauth2.is_auth_error(error) is True + + def test_detects_authentication_failed(self): + """Test detects 'authentication failed' error.""" + error = Exception("Authentication failed for user@example.com") + assert imap_oauth2.is_auth_error(error) is True + + def test_case_insensitive_not_authenticated(self): + """Test case-insensitive matching for not authenticated.""" + error = Exception("NOT AUTHENTICATED") + assert imap_oauth2.is_auth_error(error) is True + + def test_case_insensitive_authentication_failed(self): + """Test case-insensitive matching for authentication failed.""" + error = Exception("AUTHENTICATION FAILED") + assert imap_oauth2.is_auth_error(error) is True + + def test_false_for_connection_error(self): + """Test returns False for connection errors.""" + error = Exception("Connection refused") + assert imap_oauth2.is_auth_error(error) is False + + def test_false_for_timeout_error(self): + """Test returns False for timeout errors.""" + error = Exception("The read operation timed out") + assert imap_oauth2.is_auth_error(error) is False + + def test_false_for_generic_error(self): + """Test returns False for generic errors.""" + error = Exception("Something went wrong") + assert imap_oauth2.is_auth_error(error) is False + + def test_false_for_folder_not_found(self): + """Test returns False for folder errors.""" + error = Exception("Folder not found: INBOX") + assert imap_oauth2.is_auth_error(error) is False diff --git a/test/test_imap_session.py b/test/test_imap_session.py new file mode 100644 index 0000000..3f286c4 --- /dev/null +++ b/test/test_imap_session.py @@ -0,0 +1,268 @@ +""" +Tests for imap_session.py + +Tests cover: +- ensure_connection() behavior with and without OAuth2 +- ensure_folder_session() connection change detection +- Folder selection and re-selection logic +- Error handling for connection and folder failures +""" + +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) + +import imap_session + + +class TestEnsureConnection: + """Tests for ensure_connection function.""" + + def test_with_oauth2_calls_refresh(self): + """Test that OAuth2 token refresh is called when oauth2 config is present.""" + conf = { + "host": "imap.example.com", + "user": "user@example.com", + "password": "pass", + "oauth2_token": "old_token", + "oauth2": { + "provider": "microsoft", + "client_id": "client-id", + "email": "user@example.com", + }, + } + mock_conn = MagicMock() + + with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh: + with patch.object( + imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn + ) as mock_ensure: + result = imap_session.ensure_connection(mock_conn, conf) + + mock_refresh.assert_called_once_with(conf, "old_token") + mock_ensure.assert_called_once_with(mock_conn, conf) + assert result is mock_conn + + def test_without_oauth2_skips_refresh(self): + """Test that token refresh is skipped when oauth2 config is not present.""" + conf = { + "host": "imap.example.com", + "user": "user@example.com", + "password": "pass", + "oauth2": None, + } + mock_conn = MagicMock() + + with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh: + with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn): + imap_session.ensure_connection(mock_conn, conf) + + mock_refresh.assert_not_called() + + def test_without_oauth2_key_skips_refresh(self): + """Test that token refresh is skipped when oauth2 key is missing entirely.""" + conf = { + "host": "imap.example.com", + "user": "user@example.com", + "password": "pass", + } + mock_conn = MagicMock() + + with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh: + with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn): + imap_session.ensure_connection(mock_conn, conf) + + mock_refresh.assert_not_called() + + def test_returns_connection_from_ensure_connection_from_conf(self): + """Test returns the connection from ensure_connection_from_conf.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + old_conn = MagicMock() + new_conn = MagicMock() + + with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn): + result = imap_session.ensure_connection(old_conn, conf) + + assert result is new_conn + + def test_returns_none_when_connection_fails(self): + """Test returns None when ensure_connection_from_conf fails.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + + with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=None): + result = imap_session.ensure_connection(MagicMock(), conf) + + assert result is None + + def test_with_none_connection(self): + """Test works correctly when passed None as connection.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + new_conn = MagicMock() + + with patch.object( + imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn + ) as mock_ensure: + result = imap_session.ensure_connection(None, conf) + + mock_ensure.assert_called_once_with(None, conf) + assert result is new_conn + + +class TestEnsureFolderSession: + """Tests for ensure_folder_session function.""" + + def test_same_connection_no_folder_select(self): + """Test that folder is not re-selected when connection stays the same.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + mock_conn = MagicMock() + + with patch.object(imap_session, "ensure_connection", return_value=mock_conn): + result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True) + + assert result_conn is mock_conn + assert success is True + mock_conn.select.assert_not_called() + + def test_new_connection_selects_folder(self): + """Test that folder is selected when connection changes.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + old_conn = MagicMock() + new_conn = MagicMock() + + with patch.object(imap_session, "ensure_connection", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True) + + assert result_conn is new_conn + assert success is True + new_conn.select.assert_called_once_with('"INBOX"', readonly=True) + + def test_new_connection_selects_folder_readwrite(self): + """Test that folder is selected with readonly=False when specified.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + old_conn = MagicMock() + new_conn = MagicMock() + + with patch.object(imap_session, "ensure_connection", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session( + old_conn, conf, "[Gmail]/All Mail", readonly=False + ) + + assert result_conn is new_conn + assert success is True + new_conn.select.assert_called_once_with('"[Gmail]/All Mail"', readonly=False) + + def test_none_connection_returns_failure(self): + """Test returns (None, False) when ensure_connection returns None.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + mock_conn = MagicMock() + + with patch.object(imap_session, "ensure_connection", return_value=None): + result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True) + + assert result_conn is None + assert success is False + + def test_folder_select_failure_returns_false(self): + """Test returns (conn, False) when folder selection fails.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + old_conn = MagicMock() + new_conn = MagicMock() + new_conn.select.side_effect = Exception("Folder not found") + + with patch.object(imap_session, "ensure_connection", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "NonExistent", readonly=True) + + assert result_conn is new_conn + assert success is False + + def test_none_old_connection_selects_folder(self): + """Test that folder is selected when old connection was None.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + new_conn = MagicMock() + + with patch.object(imap_session, "ensure_connection", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session(None, conf, "INBOX", readonly=True) + + assert result_conn is new_conn + assert success is True + new_conn.select.assert_called_once_with('"INBOX"', readonly=True) + + def test_folder_name_with_spaces(self): + """Test folder selection with folder names containing spaces.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + old_conn = MagicMock() + new_conn = MagicMock() + + with patch.object(imap_session, "ensure_connection", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "My Folder", readonly=True) + + assert success is True + new_conn.select.assert_called_once_with('"My Folder"', readonly=True) + + def test_folder_select_imap_error(self): + """Test handles IMAP-specific errors during folder selection.""" + conf = {"host": "imap.example.com", "user": "user", "password": "pass"} + old_conn = MagicMock() + new_conn = MagicMock() + new_conn.select.side_effect = Exception("IMAP error: NO SELECT failed") + + with patch.object(imap_session, "ensure_connection", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True) + + assert result_conn is new_conn + assert success is False + + +class TestEnsureFolderSessionWithOAuth2: + """Integration-style tests for ensure_folder_session with OAuth2.""" + + def test_oauth2_refresh_triggers_folder_reselect(self): + """Test that OAuth2 token refresh (new connection) triggers folder reselection.""" + conf = { + "host": "outlook.office365.com", + "user": "user@example.com", + "password": "oauth2_token_here", + "oauth2_token": "old_token", + "oauth2": { + "provider": "microsoft", + "client_id": "client-id", + "email": "user@example.com", + }, + } + old_conn = MagicMock() + new_conn = MagicMock() + + # Simulate token refresh causing a new connection + with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"): + with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn): + result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True) + + assert result_conn is new_conn + assert success is True + new_conn.select.assert_called_once_with('"INBOX"', readonly=True) + + def test_healthy_connection_no_reselect(self): + """Test that healthy connection (same object) doesn't reselect folder.""" + conf = { + "host": "outlook.office365.com", + "user": "user@example.com", + "password": "oauth2_token_here", + "oauth2_token": "valid_token", + "oauth2": { + "provider": "microsoft", + "client_id": "client-id", + "email": "user@example.com", + }, + } + mock_conn = MagicMock() + + # Simulate healthy connection (same object returned) + with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"): + with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn): + result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True) + + assert result_conn is mock_conn + assert success is True + mock_conn.select.assert_not_called()