From dd5251f55da1001460c28ab68684e1f4058828a5 Mon Sep 17 00:00:00 2001 From: Javier Callico Date: Sat, 7 Feb 2026 19:29:21 -0500 Subject: [PATCH] Enhance migrate script with restore cache, plus cleanup for code and tests (#23) * Implement incremental migration with caching support and add tests for cache behavior. Cleanup for duplicated logic. * Rewrite tests for main scripts to use end to end testing with data setup instead of mocking. * Fix linting issues. * Fix SEARCH response handling in MockIMAPHandler to return empty response when no results are found * Refactor SEARCH command handling in MockIMAPHandler to streamline response generation for ALL queries and remove redundant code. * Add cache directory creation in load_progress_cache with logging for failures * Improve cache hit detection in process_single_uid to handle locking for existing destination message IDs * Add tests for imap_common functions and enhance email parsing in restore_imap_emails * Refactor test_backup_folder_discovery to create a parent directory for unreadable path * Add tests for Gmail mode fallback folder and dest-delete functionality in restore_imap_emails * Add tests for cache behavior and error handling in migration process --- README.md | 21 +- src/backup_imap_emails.py | 9 +- src/compare_imap_folders.py | 57 +--- src/count_imap_emails.py | 50 +--- src/imap_common.py | 123 +++++++++ src/migrate_imap_emails.py | 184 ++++++++++++- src/restore_imap_emails.py | 151 ++++++----- test/test_backup_imap_emails.py | 304 ++++----------------- test/test_compare_imap_folders.py | 110 -------- test/test_count_imap_emails.py | 215 +++++++++------ test/test_migrate_imap_emails.py | 222 +++++++++++++++ test/test_migrate_with_cache.py | 248 +++++++++++++++++ test/test_restore_imap_emails.py | 431 +++++++++++++++--------------- tools/mock_imap_server.py | 85 +++--- 14 files changed, 1306 insertions(+), 904 deletions(-) create mode 100644 test/test_migrate_with_cache.py diff --git a/README.md b/README.md index d9d0c7a..a8abd50 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ This repository contains a set of Python scripts designed to migrate emails betw - **Cleanup**: Optionally deletes messages from the source after successful transfer (effectively a "Move" operation). - *Improved for Gmail*: Automatically detects "Trash" folders to ensure emails are properly binned rather than just archived. - **Sync Mode**: Optionally deletes emails from destination that no longer exist in source (`--dest-delete`). + - **Incremental Cache**: Supports `--migrate-cache ` to store a local map of processed emails, drastically speeding up subsequent runs by skipping known messages. - **Configurable**: Adjustable concurrency and batch sizes to respect server rate limits. 2. **`compare_imap_folders.py`** (The Validator) @@ -258,11 +259,23 @@ python3 migrate_imap_emails.py \ --dest-pass "dest-app-password" ``` -Notes: -- `--preserve-flags` is enabled automatically in `--gmail-mode`. -- `--dest-delete` is not supported in `--gmail-mode`. +### 1b. Fast Incremental Migration (Cached) +Use a local cache file to remember processed emails. Use this for large migrations that may be interrupted or need to run multiple times. -### 1b. Preserve Flags (Any IMAP Server) +```bash +python3 migrate_imap_emails.py \ + --src-host "imap.source.com" \ + --src-user "source" \ + --src-pass "pass" \ + --dest-host "imap.dest.com" \ + --dest-user "dest" \ + --dest-pass "pass" \ + --migrate-cache "./migration_cache" +``` + +To force a re-check of cached items without clearing the cache logic entirely, add `--full-migrate`. + +### 1c. Preserve Flags (Any IMAP Server) Preserve IMAP flags (`\Seen`, `\Flagged`, `\Answered`, `\Draft`) during migration. If an email already exists on the destination (duplicate), the script can still sync missing flags on the existing message. diff --git a/src/backup_imap_emails.py b/src/backup_imap_emails.py index 5e4de6e..bbe9200 100644 --- a/src/backup_imap_emails.py +++ b/src/backup_imap_emails.py @@ -69,14 +69,7 @@ MANIFEST_FILENAME = "labels_manifest.json" # Thread-local storage thread_local = threading.local() -print_lock = threading.Lock() - - -def safe_print(message): - t_name = threading.current_thread().name - short_name = t_name.replace("ThreadPoolExecutor-", "T-").replace("MainThread", "MAIN") - with print_lock: - print(f"[{short_name}] {message}") +safe_print = imap_common.safe_print def get_thread_connection(src_conf): diff --git a/src/compare_imap_folders.py b/src/compare_imap_folders.py index 51416d1..603f806 100644 --- a/src/compare_imap_folders.py +++ b/src/compare_imap_folders.py @@ -60,7 +60,6 @@ Examples: import argparse import os import sys -from typing import Optional import imap_common import imap_oauth2 @@ -86,56 +85,6 @@ def get_email_count(conn, folder_name): return None -def _is_ignored_local_dir(dirname: str) -> bool: - return dirname.startswith(".") or dirname == "__pycache__" - - -def list_local_folders(local_root: str) -> list[str]: - """List all folders under a local backup root in IMAP-style names. - - The local backup format is expected to mirror IMAP folder hierarchy using - subdirectories (e.g. "[Gmail]/All Mail" becomes "[Gmail]/All Mail/"). - """ - folders: set[str] = set() - - for dirpath, dirnames, _filenames in os.walk(local_root): - dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)] - - if os.path.abspath(dirpath) == os.path.abspath(local_root): - continue - - rel = os.path.relpath(dirpath, local_root) - if rel == ".": - continue - - parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)] - if not parts: - continue - - folders.add("/".join(parts)) - - return sorted(folders) - - -def get_local_email_count(local_root: str, folder_name: str) -> Optional[int]: - """Return the count of .eml files in a local folder, or None if missing/unreadable.""" - folder_path = os.path.join(local_root, *folder_name.split("/")) - if not os.path.isdir(folder_path): - return None - - try: - count = 0 - for filename in os.listdir(folder_path): - if not filename.endswith(".eml"): - continue - full_path = os.path.join(folder_path, filename) - if os.path.isfile(full_path): - count += 1 - return count - except OSError: - return None - - def main(): default_src_path = os.getenv("SRC_LOCAL_PATH") default_dest_path = os.getenv("DEST_LOCAL_PATH") @@ -333,7 +282,7 @@ def main(): # List Source Folders print("Listing folders in Source...") if src_is_local: - folders = list_local_folders(args.src_path) + folders = imap_common.list_local_folders(args.src_path) else: folders = imap_common.list_selectable_folders(src) @@ -354,12 +303,12 @@ def main(): for folder_name in folders: # Get Counts if src_is_local: - src_count = get_local_email_count(args.src_path, folder_name) + src_count = imap_common.get_local_email_count(args.src_path, folder_name) else: src_count = get_email_count(src, folder_name) if dest_is_local: - dest_count = get_local_email_count(args.dest_path, folder_name) + dest_count = imap_common.get_local_email_count(args.dest_path, folder_name) else: dest_count = get_email_count(dest, folder_name) diff --git a/src/count_imap_emails.py b/src/count_imap_emails.py index 53fc95a..66353ad 100644 --- a/src/count_imap_emails.py +++ b/src/count_imap_emails.py @@ -109,56 +109,10 @@ def count_emails(imap_server, username, password=None, oauth2_token=None): print(f"An error occurred: {e}") -def _is_ignored_local_dir(dirname: str) -> bool: - return dirname.startswith(".") or dirname == "__pycache__" - - -def list_local_folders(local_root: str) -> list[str]: - """List all folders under a local backup root in IMAP-style names.""" - folders: set[str] = set() - - for dirpath, dirnames, _filenames in os.walk(local_root): - dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)] - - if os.path.abspath(dirpath) == os.path.abspath(local_root): - continue - - rel = os.path.relpath(dirpath, local_root) - if rel == ".": - continue - - parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)] - if not parts: - continue - - folders.add("/".join(parts)) - - return sorted(folders) - - -def get_local_email_count(local_root: str, folder_name: str): - """Return the count of .eml files in a local folder, or None if missing/unreadable.""" - folder_path = os.path.join(local_root, *folder_name.split("/")) - if not os.path.isdir(folder_path): - return None - - try: - count = 0 - for filename in os.listdir(folder_path): - if not filename.endswith(".eml"): - continue - full_path = os.path.join(folder_path, filename) - if os.path.isfile(full_path): - count += 1 - return count - except OSError: - return None - - def count_local_emails(local_path: str) -> None: print(f"Scanning local backup: {local_path}") - folders = list_local_folders(local_path) + folders = imap_common.list_local_folders(local_path) if not folders: print("No folders found.") return @@ -168,7 +122,7 @@ def count_local_emails(local_path: str) -> None: print("-" * 52) for folder_name in folders: - count = get_local_email_count(local_path, folder_name) + count = imap_common.get_local_email_count(local_path, folder_name) if count is None: print(f"{folder_name:<40} {'N/A':>10}") continue diff --git a/src/imap_common.py b/src/imap_common.py index bde653a..53a96f8 100644 --- a/src/imap_common.py +++ b/src/imap_common.py @@ -10,11 +10,13 @@ import imaplib import os import re import sys +import threading from email import policy from email.header import decode_header from email.parser import BytesParser import imap_oauth2 +import restore_cache # Standard IMAP flags FLAG_SEEN = "\\Seen" @@ -42,6 +44,127 @@ CMD_SEARCH = "search" CMD_FETCH = "fetch" OP_ADD_FLAGS = "+FLAGS" +_print_lock = threading.Lock() + + +def safe_print(message: str) -> None: + """Thread-safe print with short thread names for logs.""" + t_name = threading.current_thread().name + short_name = t_name.replace("ThreadPoolExecutor-", "T-").replace("MainThread", "MAIN") + with _print_lock: + print(f"[{short_name}] {message}") + + +def _is_ignored_local_dir(dirname: str) -> bool: + return dirname.startswith(".") or dirname == "__pycache__" + + +def list_local_folders(local_root: str) -> list[str]: + """List all folders under a local backup root in IMAP-style names.""" + folders: set[str] = set() + + for dirpath, dirnames, _filenames in os.walk(local_root): + dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)] + + if os.path.abspath(dirpath) == os.path.abspath(local_root): + continue + + rel = os.path.relpath(dirpath, local_root) + if rel == ".": + continue + + parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)] + if not parts: + continue + + folders.add("/".join(parts)) + + return sorted(folders) + + +def get_local_email_count(local_root: str, folder_name: str) -> int | None: + """Return the count of .eml files in a local folder, or None if missing/unreadable.""" + folder_path = os.path.join(local_root, *folder_name.split("/")) + if not os.path.isdir(folder_path): + return None + + try: + count = 0 + for filename in os.listdir(folder_path): + if not filename.endswith(".eml"): + continue + full_path = os.path.join(folder_path, filename) + if os.path.isfile(full_path): + count += 1 + return count + except OSError: + return None + + +def get_backup_folders(local_path: str) -> list[tuple[str, str]]: + """Scan the backup directory and return list of (folder_name, folder_path) tuples.""" + folders: list[tuple[str, str]] = [] + + def scan_dir(path: str, prefix: str = "") -> None: + try: + for item in os.listdir(path): + item_path = os.path.join(path, item) + if os.path.isdir(item_path): + # Check if this directory contains .eml files + has_eml = any( + f.endswith(".eml") for f in os.listdir(item_path) if os.path.isfile(os.path.join(item_path, f)) + ) + folder_name = f"{prefix}{item}" if prefix else item + + if has_eml: + folders.append((folder_name, item_path)) + + # Recurse into subdirectories + scan_dir(item_path, f"{folder_name}/") + except Exception: + pass + + scan_dir(local_path) + return folders + + +def extract_message_id_from_eml(file_path: str, read_limit: int = 65536) -> str | None: + """Extract just the Message-ID from an .eml file efficiently.""" + try: + with open(file_path, "rb") as f: + header_bytes = f.read(read_limit) + + return extract_message_id(header_bytes) + except Exception: + return None + + +def is_progress_cache_ready(cache_data: dict | None, cache_lock: threading.Lock | None) -> bool: + """Return True when cache data and lock are initialized.""" + return cache_data is not None and cache_lock is not None + + +def load_progress_cache( + cache_root: str, + dest_host: str, + dest_user: str, + *, + log_fn=None, +) -> tuple[str, dict, threading.Lock]: + """Load or initialize a progress cache file for a destination.""" + try: + os.makedirs(cache_root, exist_ok=True) + except Exception as exc: + if log_fn is not None: + log_fn(f"Warning: unable to create cache directory '{cache_root}': {exc}") + + cache_path = restore_cache.get_dest_index_cache_path(cache_root, dest_host, dest_user) + cache_data = restore_cache.load_dest_index_cache(cache_path) + cache_lock = threading.Lock() + if log_fn is not None: + log_fn(f"Using progress cache: {cache_path}") + return cache_path, cache_data, cache_lock + def ensure_folder_exists(imap_conn, folder_name: str) -> None: """Best-effort create of a folder if it doesn't already exist. diff --git a/src/migrate_imap_emails.py b/src/migrate_imap_emails.py index 27afeba..a626a6f 100644 --- a/src/migrate_imap_emails.py +++ b/src/migrate_imap_emails.py @@ -16,6 +16,10 @@ Features: applies additional Gmail labels by copying the message into label folders. - In Gmail mode, label preservation is enabled automatically. - Note: --dest-delete is not supported in --gmail-mode. +- Cached Incremental Migration (--migrate-cache): + - Uses a local JSON cache to track migrated Message-IDs. + - Dramatically speeds up re-runs by skipping already processed emails without server checks. + - Use --full-migrate to ignore cache skipping (force check) while still updating cache. Configuration (Environment Variables): Source Account: @@ -106,6 +110,17 @@ Usage Example: --dest-host "imap.gmail.com" \ --dest-user "dest@gmail.com" \ --dest-pass "DEST_APP_PASSWORD" + + # Cached Incremental Migration (Recommended for large accounts): + # Uses a local cache to track progress and skip already migrated emails. + python3 migrate_imap_emails.py \ + --migrate-cache "./migration_cache" \ + --src-host "imap.example.com" \ + --src-user "source@example.com" \ + --src-pass "SOURCE_PASSWORD" \ + --dest-host "imap.example.com" \ + --dest-user "dest@example.com" \ + --dest-pass "DEST_PASSWORD" """ import argparse @@ -114,12 +129,14 @@ import os import re import sys import threading +from typing import Optional import imap_common import imap_oauth2 import imap_session import provider_exchange import provider_gmail +import restore_cache # Configuration defaults DELETE_FROM_SOURCE_DEFAULT = False @@ -128,15 +145,7 @@ BATCH_SIZE = 10 # Initial default, updated in main # Thread-local storage for IMAP connections thread_local = threading.local() -print_lock = threading.Lock() - - -def safe_print(message): - t_name = threading.current_thread().name - # Shorten thread name for cleaner logs e.g. ThreadPoolExecutor-0_0 -> T-0_0 - short_name = t_name.replace("ThreadPoolExecutor-", "T-").replace("MainThread", "MAIN") - with print_lock: - print(f"[{short_name}] {message}") +safe_print = imap_common.safe_print def filter_preservable_flags(flags_str): @@ -255,6 +264,14 @@ def process_single_uid( gmail_mode, label_index, check_duplicate, + full_migrate: bool = False, + existing_dest_msg_ids: Optional[set[str]] = None, + existing_dest_msg_ids_lock: Optional[threading.Lock] = None, + progress_cache_path: Optional[str] = None, + progress_cache_data: Optional[dict] = None, + progress_cache_lock: Optional[threading.Lock] = None, + dest_host: Optional[str] = None, + dest_user: Optional[str] = None, ): """ Migrate a single email by UID. @@ -327,10 +344,27 @@ def process_single_uid( 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)) + is_cached = False + is_duplicate = False + + # check cache first if available + cache_hit = False + if not full_migrate and msg_id and existing_dest_msg_ids is not None: + if existing_dest_msg_ids_lock is not None: + with existing_dest_msg_ids_lock: + cache_hit = msg_id in existing_dest_msg_ids + else: + cache_hit = msg_id in existing_dest_msg_ids + + if cache_hit: + is_cached = True + is_duplicate = True + safe_print(f"[{target_folder}] SKIP (cached) | {size_str:<8} | {subject[:40]}") + elif bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id)): + is_duplicate = True + safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}") 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: @@ -348,6 +382,21 @@ def process_single_uid( for flag in flags.split(): safe_print(f" -> Applied flag: {flag}") + # Update cache if processed effectively (copied or duplicate) + if msg_id: + restore_cache.record_progress( + message_id=msg_id, + folder_name=folder_name, + existing_dest_msg_ids=existing_dest_msg_ids, + existing_dest_msg_ids_lock=existing_dest_msg_ids_lock, + progress_cache_path=progress_cache_path, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, + dest_host=dest_host, + dest_user=dest_user, + log_fn=safe_print, + ) + # Apply remaining Gmail labels if apply_labels and remaining_labels and msg_id: for label in remaining_labels: @@ -414,6 +463,12 @@ def process_batch( gmail_mode=False, label_index=None, check_duplicate=True, + full_migrate=False, + existing_dest_msg_ids: Optional[set[str]] = None, + existing_dest_msg_ids_lock: Optional[threading.Lock] = None, + progress_cache_path: Optional[str] = None, + progress_cache_data: Optional[dict] = None, + progress_cache_lock: Optional[threading.Lock] = None, ): src, dest = get_thread_connections(src_conf, dest_conf) if not src or not dest: @@ -434,6 +489,10 @@ def process_batch( safe_print(f"Error selecting folder {folder_name} in worker: {e}") return + # Extract info for cache update if needed + dest_host = dest_conf.get("host") + dest_user = dest_conf.get("user") + deleted_count = 0 for uid in uids: @@ -471,6 +530,14 @@ def process_batch( gmail_mode, label_index, check_duplicate, + full_migrate, + existing_dest_msg_ids=existing_dest_msg_ids, + existing_dest_msg_ids_lock=existing_dest_msg_ids_lock, + progress_cache_path=progress_cache_path, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, + dest_host=dest_host, + dest_user=dest_user, ) thread_local.src = src thread_local.dest = dest @@ -504,9 +571,49 @@ def migrate_folder( preserve_flags=False, gmail_mode=False, label_index=None, + progress_cache_path: Optional[str] = None, + full_migrate=False, + progress_cache_file: Optional[str] = None, + progress_cache_data: Optional[dict] = None, + progress_cache_lock: Optional[threading.Lock] = None, ): safe_print(f"--- Preparing Folder: {folder_name} ---") + # Load cache if provided + existing_dest_msg_ids = None + existing_dest_msg_ids_lock = None + dest_host = dest_conf.get("host") + dest_user = dest_conf.get("user") + cache_file = progress_cache_file + + if progress_cache_path: + if progress_cache_data is None or progress_cache_lock is None or cache_file is None: + try: + cache_file, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache( + progress_cache_path, + dest_host, + dest_user, + log_fn=safe_print, + ) + except Exception as e: + safe_print(f"Warning: Failed to load cache: {e}") + + if imap_common.is_progress_cache_ready(progress_cache_data, progress_cache_lock): + try: + existing_dest_msg_ids = restore_cache.get_cached_message_ids( + progress_cache_data, + progress_cache_lock, + dest_host, + dest_user, + folder_name, + ) + existing_dest_msg_ids_lock = threading.Lock() + safe_print(f"Cache has {len(existing_dest_msg_ids)} Message-IDs for this folder.") + except Exception as e: + safe_print(f"Warning: Failed to read cache for folder '{folder_name}': {e}") + existing_dest_msg_ids = set() + existing_dest_msg_ids_lock = threading.Lock() + # Maintain folder structure (skip in Gmail mode; worker will create/select target label folders) if not gmail_mode: try: @@ -547,7 +654,9 @@ def migrate_folder( safe_print(f"Pre-fetching destination Message-IDs for {folder_name}...") dest_uid_to_msgid = imap_common.get_message_ids_in_folder(dest) dest_msg_ids = set(dest_uid_to_msgid.values()) - safe_print(f"Found {len(dest_msg_ids)} existing messages in destination.") + if existing_dest_msg_ids and not full_migrate: + dest_msg_ids.update(existing_dest_msg_ids) + safe_print(f"Found {len(dest_msg_ids)} existing messages in destination (server + cache).") # Pre-fetch source Message-IDs and filter out duplicates before processing # Skip pre-filtering when preserve_flags is True (need to sync flags on duplicates) @@ -603,6 +712,12 @@ def migrate_folder( gmail_mode, label_index, check_duplicate, + full_migrate, + existing_dest_msg_ids, + existing_dest_msg_ids_lock, + cache_file, + progress_cache_data, + progress_cache_lock, ) ) @@ -642,6 +757,10 @@ def migrate_folder( safe_print("Syncing destination: removing emails not in source...") delete_orphan_emails(dest, folder_name, src_msg_ids, dest_uid_to_msgid) + # Force-flush progress cache at end of folder migration. + if cache_file and imap_common.is_progress_cache_ready(progress_cache_data, progress_cache_lock): + restore_cache.maybe_save_dest_index_cache(cache_file, progress_cache_data, progress_cache_lock, force=True) + def main(): parser = argparse.ArgumentParser(description="Migrate emails between IMAP accounts.") @@ -781,6 +900,16 @@ def main(): help="Gmail migration mode", ) + parser.add_argument( + "--migrate-cache", + help="Path to directory for migration progress cache (enables incremental migration)", + ) + parser.add_argument( + "--full-migrate", + action="store_true", + help="Force full migration (ignore cache for skipping), but still update cache if --migrate-cache provided", + ) + args = parser.parse_args() # Assign to variables @@ -794,6 +923,8 @@ def main(): DEST_DELETE = args.dest_delete gmail_mode = bool(args.gmail_mode) + migrate_cache = args.migrate_cache + full_migrate = args.full_migrate preserve_flags = bool(args.preserve_flags) or gmail_mode preserve_labels = bool(args.preserve_labels) or gmail_mode @@ -896,6 +1027,20 @@ def main(): else None, } + progress_cache_file = None + progress_cache_data = None + progress_cache_lock = None + if migrate_cache: + try: + progress_cache_file, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache( + migrate_cache, + DEST_HOST, + DEST_USER, + log_fn=safe_print, + ) + except Exception as e: + safe_print(f"Warning: Failed to load progress cache: {e}") + try: # Initial connection to list folders safe_print("Connecting to Source to list folders...") @@ -949,6 +1094,11 @@ def main(): preserve_flags, gmail_mode, label_index, + progress_cache_path=migrate_cache, + full_migrate=full_migrate, + progress_cache_file=progress_cache_file, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, ) else: # Migration for all folders @@ -972,6 +1122,11 @@ def main(): preserve_flags, True, label_index, + progress_cache_path=migrate_cache, + full_migrate=full_migrate, + progress_cache_file=progress_cache_file, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, ) if not gmail_mode: @@ -1005,6 +1160,11 @@ def main(): preserve_flags, False, None, + progress_cache_path=migrate_cache, + full_migrate=full_migrate, + progress_cache_file=progress_cache_file, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, ) src_main.logout() diff --git a/src/restore_imap_emails.py b/src/restore_imap_emails.py index 84e43b3..8716407 100644 --- a/src/restore_imap_emails.py +++ b/src/restore_imap_emails.py @@ -80,14 +80,7 @@ BATCH_SIZE = 10 # Thread-local storage thread_local = threading.local() -print_lock = threading.Lock() - - -def safe_print(message): - t_name = threading.current_thread().name - short_name = t_name.replace("ThreadPoolExecutor-", "T-").replace("MainThread", "MAIN") - with print_lock: - print(f"[{short_name}] {message}") +safe_print = imap_common.safe_print def get_thread_connection(dest_conf): @@ -209,22 +202,6 @@ def sync_flags_on_existing(imap_conn, folder_name, message_id, flags, size): safe_print(f" -> Error syncing flags: {e}") -def extract_message_id_from_eml(file_path): - """ - Extract just the Message-ID from an .eml file efficiently. - Returns the Message-ID string or None on error. - """ - try: - with open(file_path, "rb") as f: - # Read just the first 64KB to get headers - header_bytes = f.read(65536) - - msg_id = imap_common.extract_message_id(header_bytes) - return msg_id - except Exception: - return None - - def parse_eml_file(file_path): """ Parse an .eml file and extract metadata. @@ -453,7 +430,11 @@ def process_restore_batch( # Lazy-load per-folder progress cache (no destination scan). if existing_dest_msg_ids is None: built: set[str] = set() - if progress_cache_data is not None and progress_cache_lock is not None and dest_host and dest_user: + if ( + imap_common.is_progress_cache_ready(progress_cache_data, progress_cache_lock) + and dest_host + and dest_user + ): built = restore_cache.get_cached_message_ids( progress_cache_data, progress_cache_lock, @@ -555,8 +536,10 @@ def process_restore_batch( if label_folder_msg_ids is None: built: set[str] = set() if ( - progress_cache_data is not None - and progress_cache_lock is not None + imap_common.is_progress_cache_ready( + progress_cache_data, + progress_cache_lock, + ) and dest_host and dest_user ): @@ -689,6 +672,9 @@ def restore_folder( dest_delete=False, full_restore: bool = False, cache_root: Optional[str] = None, + progress_cache_file: Optional[str] = None, + progress_cache_data: Optional[dict] = None, + progress_cache_lock: Optional[threading.Lock] = None, ): """ Restore all emails from a local folder to the destination IMAP server. @@ -709,11 +695,16 @@ def restore_folder( safe_print(f"Found {len(eml_files)} emails to restore.") cache_root = cache_root or local_folder_path - cache_path = restore_cache.get_dest_index_cache_path(cache_root, dest_conf["host"], dest_conf["user"]) - cache_data: dict = restore_cache.load_dest_index_cache(cache_path) - cache_lock = threading.Lock() - - safe_print(f"Using progress cache: {cache_path}") + cache_path = progress_cache_file + cache_data = progress_cache_data + cache_lock = progress_cache_lock + if cache_path is None or cache_data is None or cache_lock is None: + cache_path, cache_data, cache_lock = imap_common.load_progress_cache( + cache_root, + dest_conf["host"], + dest_conf["user"], + log_fn=safe_print, + ) # Incremental mode uses cached Message-IDs to skip already-processed emails. existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {folder_name: set()} @@ -768,7 +759,7 @@ def restore_folder( files_to_restore = [] skipped = 0 for file_path, filename in eml_files: - msg_id = extract_message_id_from_eml(file_path) + msg_id = imap_common.extract_message_id_from_eml(file_path) if msg_id and msg_id in dest_msg_ids: skipped += 1 else: @@ -777,6 +768,14 @@ def restore_folder( if not files_to_restore: safe_print("No new emails to restore.") + if dest_delete and local_msg_ids is not None: + safe_print("Syncing destination: removing emails not in local backup...") + dest = imap_session.ensure_connection(None, dest_conf) + if dest: + delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids) + dest.logout() + + restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True) return safe_print(f"Starting parallel restore of {len(files_to_restore)} emails...") @@ -825,7 +824,16 @@ def restore_folder( restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True) -def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore: bool = False): +def restore_gmail_with_labels( + local_path, + dest_conf, + manifest, + apply_flags, + full_restore: bool = False, + progress_cache_file: Optional[str] = None, + progress_cache_data: Optional[dict] = None, + progress_cache_lock: Optional[threading.Lock] = None, +): """ Special restoration mode for Gmail: Upload emails to their first label folder and then apply additional labels from the manifest. @@ -856,11 +864,14 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)] - cache_path = restore_cache.get_dest_index_cache_path(local_path, dest_conf["host"], dest_conf["user"]) - progress_cache_data: dict = restore_cache.load_dest_index_cache(cache_path) - progress_cache_lock = threading.Lock() - - safe_print(f"Using progress cache: {cache_path}") + cache_path = progress_cache_file + if cache_path is None or progress_cache_data is None or progress_cache_lock is None: + cache_path, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache( + local_path, + dest_conf["host"], + dest_conf["user"], + log_fn=safe_print, + ) safe_print( "Cache will be populated as restore runs (no up-front destination indexing). " "First run may still do per-message duplicate checks; subsequent runs will skip quickly." @@ -913,36 +924,6 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full restore_cache.maybe_save_dest_index_cache(cache_path, progress_cache_data, progress_cache_lock, force=True) -def get_backup_folders(local_path): - """ - Scan the backup directory and return list of folder paths. - Returns list of (folder_name, local_path) tuples. - """ - folders = [] - - def scan_dir(path, prefix=""): - try: - for item in os.listdir(path): - item_path = os.path.join(path, item) - if os.path.isdir(item_path): - # Check if this directory contains .eml files - has_eml = any( - f.endswith(".eml") for f in os.listdir(item_path) if os.path.isfile(os.path.join(item_path, f)) - ) - folder_name = f"{prefix}{item}" if prefix else item - - if has_eml: - folders.append((folder_name, item_path)) - - # Recurse into subdirectories - scan_dir(item_path, f"{folder_name}/") - except Exception: - pass - - scan_dir(local_path) - return folders - - def main(): parser = argparse.ArgumentParser(description="Restore IMAP emails from local .eml files.") @@ -1126,6 +1107,19 @@ def main(): if not manifest: print("Warning: No manifest found. Labels/flags will not be applied.") + progress_cache_file = None + progress_cache_data = None + progress_cache_lock = None + try: + progress_cache_file, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache( + local_path, + args.dest_host, + args.dest_user, + log_fn=safe_print, + ) + except Exception as e: + safe_print(f"Warning: Failed to load progress cache: {e}") + print("\n--- Configuration Summary ---") print(f"Source Path : {local_path}") print(f"Destination Host: {args.dest_host}") @@ -1159,7 +1153,16 @@ def main(): if args.gmail_mode: dest.logout() # Special Gmail mode - restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore=args.full_restore) + restore_gmail_with_labels( + local_path, + dest_conf, + manifest, + apply_flags, + full_restore=args.full_restore, + progress_cache_file=progress_cache_file, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, + ) dest = None # Connection handled by restore_gmail_with_labels elif args.folder: # Restore specific folder @@ -1177,11 +1180,14 @@ def main(): args.dest_delete, full_restore=args.full_restore, cache_root=local_path, + progress_cache_file=progress_cache_file, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, ) dest.logout() else: # Restore all folders - folders = get_backup_folders(local_path) + folders = imap_common.get_backup_folders(local_path) if not folders: print("No backup folders found.") sys.exit(1) @@ -1208,6 +1214,9 @@ def main(): args.dest_delete, full_restore=args.full_restore, cache_root=local_path, + progress_cache_file=progress_cache_file, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, ) dest.logout() diff --git a/test/test_backup_imap_emails.py b/test/test_backup_imap_emails.py index 94e4f46..107d545 100644 --- a/test/test_backup_imap_emails.py +++ b/test/test_backup_imap_emails.py @@ -9,9 +9,9 @@ Tests cover: - Configuration validation """ +import imaplib import os import sys -from unittest.mock import MagicMock import pytest @@ -250,124 +250,6 @@ class TestGetExistingUids: result = backup_imap_emails.get_existing_uids(str(tmp_path)) assert result == {"1"} - def test_os_error_handling(self, monkeypatch): - """Test handling of OS errors during listing.""" - - def mock_listdir(path): - raise OSError("Access denied") - - monkeypatch.setattr(os, "listdir", mock_listdir) - - uids = backup_imap_emails.get_existing_uids("/some/path") - assert len(uids) == 0 - - -class TestBackupErrorHandling: - """Tests for error handling scenarios in backup.""" - - def test_connection_error_in_worker(self, monkeypatch, tmp_path): - """Test worker handles connection failure gracefully.""" - # Mock get_imap_connection to fail - monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: None) - - # Should return None/Exit without crashing - backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), str(tmp_path)) - - def test_select_error_in_worker(self, monkeypatch, tmp_path): - """Test worker handles SELECT failure.""" - mock_conn = MagicMock() - mock_conn.select.side_effect = Exception("Select error") - monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn) - - # Should log error and return - backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), str(tmp_path)) - mock_conn.select.assert_called() - - def test_fetch_body_error(self, monkeypatch, tmp_path): - """Test handling of fetch body failure.""" - mock_conn = MagicMock() - mock_conn.select.return_value = "OK" - - # Fetch body fails - mock_conn.uid.return_value = ("NO", [None]) - - monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn) - - # Try processing one UID - backup_imap_emails.process_batch([b"1"], "INBOX", ("h", "u", "p"), str(tmp_path)) - - # File should not exist - assert not list(tmp_path.glob("*.eml")) - - def test_write_error(self, monkeypatch, tmp_path): - """Test handling of file write error.""" - mock_conn = MagicMock() - monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn) - - # Mock fetched data - mock_conn.uid.return_value = ("OK", [(b"1 (RFC822 {10}", b"Content")]) - - # Mock open to fail - def mock_open(*args, **kwargs): - raise OSError("Disk full") - - monkeypatch.setattr("builtins.open", mock_open) - - backup_imap_emails.process_batch([b"1"], "INBOX", ("h", "u", "p"), str(tmp_path)) - - def test_folder_creation_error(self, monkeypatch, tmp_path): - """Test handling failure to create local folder.""" - - def mock_makedirs(path, exist_ok=False): - raise OSError("Permission denied") - - monkeypatch.setattr(os, "makedirs", mock_makedirs) - - mock_conn = MagicMock() - - # backup_folder should return early - backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p")) - mock_conn.select.assert_not_called() - - def test_select_folder_error(self, monkeypatch, tmp_path): - """Test handling of select folder failure in main loop.""" - mock_conn = MagicMock() - mock_conn.select.side_effect = Exception("Select failed") - monkeypatch.setattr(os, "makedirs", lambda p, exist_ok: None) - - backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p")) - mock_conn.uid.assert_not_called() - - def test_search_error(self, monkeypatch, tmp_path): - """Test handling of search failure.""" - mock_conn = MagicMock() - mock_conn.uid.return_value = ("NO", []) - monkeypatch.setattr(os, "makedirs", lambda p, exist_ok: None) - - backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p")) - - def test_main_makedirs_error(self, monkeypatch, capsys): - """Test failure to create main backup directory.""" - env = { - "SRC_IMAP_HOST": "h", - "SRC_IMAP_USERNAME": "u", - "SRC_IMAP_PASSWORD": "p", - } - monkeypatch.setattr(os, "environ", env) - monkeypatch.setattr(sys, "argv", ["backup.py", "--dest-path", "/protected/path"]) - - def mock_makedirs(path): - raise OSError("No permission") - - monkeypatch.setattr(os, "makedirs", mock_makedirs) - monkeypatch.setattr(os.path, "exists", lambda p: False) - - with pytest.raises(SystemExit): - backup_imap_emails.main() - - captured = capsys.readouterr() - assert "Error creating backup directory" in captured.out - class TestGmailLabelsPreservation: """Tests for Gmail labels manifest functionality.""" @@ -419,173 +301,77 @@ class TestGmailLabelsPreservation: result = backup_imap_emails.load_labels_manifest(str(tmp_path)) assert result == {} - def test_get_message_ids_in_folder(self, monkeypatch): - """Test extraction of message IDs from a folder.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.uid.side_effect = [ - ("OK", [b"1 2 3"]), # search result - ( - "OK", - [ - (b"1 (FLAGS (\\Seen) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: \r\n"), - b")", - (b"2 (FLAGS () BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: \r\n"), - b")", - ( - b"3 (FLAGS (\\Seen \\Answered) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", - b"Message-ID: \r\n", - ), - b")", - ], - ), # fetch result - ] + def test_get_message_ids_in_folder(self, single_mock_server): + src_data = { + "INBOX": [ + b"Subject: A\r\nMessage-ID: \r\n\r\nBody", + b"Subject: B\r\nMessage-ID: \r\n\r\nBody", + ] + } + _server, port = single_mock_server(src_data) - result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", None) + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") + + result = backup_imap_emails.get_message_ids_in_folder(conn, "INBOX", None) assert "" in result assert "" in result - assert "" in result + conn.logout() - def test_get_message_info_in_folder_read_status(self, monkeypatch): - """Test extraction of message IDs with read/unread status.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.uid.side_effect = [ - ("OK", [b"1 2 3"]), # search result - ( - "OK", - [ - (b"1 (FLAGS (\\Seen) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: \r\n"), - b")", - (b"2 (FLAGS () BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: \r\n"), - b")", - ( - b"3 (FLAGS (\\Seen \\Answered) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", - b"Message-ID: \r\n", - ), - b")", - ], - ), # fetch result - ] + def test_get_message_info_in_folder_read_status(self, single_mock_server): + src_data = { + "INBOX": [ + {"uid": 1, "flags": {"\\Seen"}, "content": b"Message-ID: \r\n"}, + {"uid": 2, "flags": set(), "content": b"Message-ID: \r\n"}, + {"uid": 3, "flags": {"\\Seen", "\\Answered"}, "content": b"Message-ID: \r\n"}, + ] + } + _server, port = single_mock_server(src_data) - result = backup_imap_emails.get_message_info_in_folder(mock_conn, "INBOX", None) + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") + + result = backup_imap_emails.get_message_info_in_folder(conn, "INBOX", None) assert "" in result - assert "\\Seen" in result[""]["flags"] # Has \Seen flag + assert "\\Seen" in result[""]["flags"] assert "" in result - assert result[""]["flags"] == [] # No flags + assert result[""]["flags"] == [] assert "" in result - assert "\\Seen" in result[""]["flags"] # Has \Seen flag - assert "\\Answered" in result[""]["flags"] # Also has \Answered + assert "\\Seen" in result[""]["flags"] + assert "\\Answered" in result[""]["flags"] + conn.logout() - def test_get_message_ids_in_folder_with_progress(self, monkeypatch): - """Test extraction of message IDs with progress callback.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.uid.side_effect = [ - ("OK", [b"1 2 3"]), # search result - ( - "OK", - [ - (b"1 (FLAGS (\\Seen) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: \r\n"), - b")", - ], - ), # fetch result - ] - - progress_calls = [] - - def progress_cb(current, total): - progress_calls.append((current, total)) - - backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", progress_cb) - - # Progress should have been called - assert len(progress_calls) > 0 - # Last call should show completion - assert progress_calls[-1][0] == progress_calls[-1][1] - - def test_get_message_ids_in_folder_select_error(self, monkeypatch): - """Test handling of folder select error.""" - mock_conn = MagicMock() - mock_conn.select.side_effect = Exception("Select failed") - - result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX") - assert result == set() - - def test_get_message_ids_in_folder_empty(self, monkeypatch): - """Test extraction from empty folder.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"0"]) - mock_conn.uid.return_value = ("OK", [b""]) - - result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", None) - assert result == set() - - def test_build_labels_manifest(self, monkeypatch, tmp_path): - """Test building labels manifest from mock folders.""" - mock_conn = MagicMock() - - # Mock folder list - simulate Gmail structure - mock_conn.list.return_value = ( - "OK", - [ - b'(\\HasNoChildren) "/" "INBOX"', - b'(\\HasNoChildren) "/" "Work"', - b'(\\HasNoChildren) "/" "[Gmail]/All Mail"', - b'(\\HasNoChildren) "/" "[Gmail]/Sent Mail"', + def test_build_labels_manifest(self, single_mock_server, tmp_path): + src_data = { + "INBOX": [b"Message-ID: \r\n"], + "Work": [b"Message-ID: \r\n"], + "[Gmail]/Sent Mail": [b"Message-ID: \r\n"], + "[Gmail]/All Mail": [ + {"uid": 1, "flags": {"\\Seen", "\\Flagged"}, "content": b"Message-ID: \r\n"}, + {"uid": 2, "flags": set(), "content": b"Message-ID: \r\n"}, ], - ) - - # Track which folder is selected - folder_data = { - "INBOX": {"", ""}, - "Work": {""}, - "[Gmail]/Sent Mail": {""}, } + _server, port = single_mock_server(src_data) - # Mock info for All Mail (with flags) - all_mail_info = { - "": {"flags": ["\\Seen", "\\Flagged"]}, - "": {"flags": []}, - } + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") - def mock_get_message_info_with_conf(conn, folder, src_conf, progress_cb=None): - if folder == "[Gmail]/All Mail": - return (all_mail_info, conn) - return ({}, conn) + result = backup_imap_emails.build_labels_manifest(conn, str(tmp_path)) - 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)) - - # Check manifest structure (new format with labels and flags) assert "" in result assert "" in result assert "INBOX" in result[""]["labels"] assert "Work" in result[""]["labels"] assert "\\Seen" in result[""]["flags"] assert "\\Flagged" in result[""]["flags"] - assert "INBOX" in result[""]["labels"] assert "Sent Mail" in result[""]["labels"] assert result[""]["flags"] == [] - # Check file was saved manifest_path = tmp_path / "labels_manifest.json" assert manifest_path.exists() - - def test_build_labels_manifest_list_error(self, monkeypatch, tmp_path): - """Test handling of folder list error.""" - mock_conn = MagicMock() - mock_conn.list.return_value = ("NO", []) - - result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path)) - assert result == {} + conn.logout() def test_preserve_labels_flag_integration(self, single_mock_server, monkeypatch, tmp_path): """Test --preserve-labels flag creates manifest.""" diff --git a/test/test_compare_imap_folders.py b/test/test_compare_imap_folders.py index 75a092a..5c73877 100644 --- a/test/test_compare_imap_folders.py +++ b/test/test_compare_imap_folders.py @@ -13,7 +13,6 @@ Tests cover: import imaplib import os import sys -from unittest.mock import MagicMock import pytest @@ -171,115 +170,6 @@ class TestGetEmailCount: conn.logout() -class TestCompareFoldersErrorHandling: - """Tests for error handling in compare_imap_folders.py""" - - def test_connection_failure(self, monkeypatch): - """Test graceful exit when connection fails.""" - mock_get = MagicMock(return_value=None) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) - - # Test source fail - env = { - "SRC_IMAP_HOST": "h", - "SRC_IMAP_USERNAME": "u", - "SRC_IMAP_PASSWORD": "p", - "DEST_IMAP_HOST": "h", - "DEST_IMAP_USERNAME": "u", - "DEST_IMAP_PASSWORD": "p", - } - monkeypatch.setattr(os, "environ", env) - monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"]) - - compare_imap_folders.main() - # Should call once and exit - assert mock_get.call_count == 1 - - def test_dest_connection_failure(self, monkeypatch): - """Test graceful exit when destination connection fails.""" - # Source OK, Dest None - mock_src = MagicMock() - mock_src.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"']) - - def side_effect(h, u, p): - if u == "src_u": - return mock_src - return None - - monkeypatch.setattr("imap_common.get_imap_connection", side_effect) - - env = { - "SRC_IMAP_HOST": "h", - "SRC_IMAP_USERNAME": "src_u", - "SRC_IMAP_PASSWORD": "p", - "DEST_IMAP_HOST": "h", - "DEST_IMAP_USERNAME": "dest_u", - "DEST_IMAP_PASSWORD": "p", - } - monkeypatch.setattr(os, "environ", env) - monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"]) - - compare_imap_folders.main() - - def test_list_failure(self, monkeypatch, capsys): - """Test handling of LIST command failure.""" - mock_src = MagicMock() - mock_src.list.return_value = ("NO", []) - mock_dest = MagicMock() - - monkeypatch.setattr( - "imap_common.get_imap_connection", lambda h, u, p, oauth2_token=None: mock_src if u == "s" else mock_dest - ) - - env = { - "SRC_IMAP_HOST": "h", - "SRC_IMAP_USERNAME": "s", - "SRC_IMAP_PASSWORD": "p", - "DEST_IMAP_HOST": "h", - "DEST_IMAP_USERNAME": "d", - "DEST_IMAP_PASSWORD": "p", - } - monkeypatch.setattr(os, "environ", env) - monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"]) - - compare_imap_folders.main() - - captured = capsys.readouterr() - assert "Failed to list source folders" in captured.out - - def test_select_failure(self, single_mock_server): - """Test get_email_count handles select failure.""" - data = {"INBOX": []} - _, port = single_mock_server(data) - conn = imaplib.IMAP4("localhost", port) - conn.login("user", "pass") - - # Mocking select failure on a real connection object is hard without - # using a pure mock. So let's use a Mock object instead of real conn. - mock_conn = MagicMock() - mock_conn.select.return_value = ("NO", [b"Error"]) - - result = compare_imap_folders.get_email_count(mock_conn, "INBOX") - assert result is None - - def test_search_failure(self, single_mock_server): - """Test get_email_count handles search failure.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"Selected"]) - mock_conn.search.return_value = ("NO", [b"Error"]) - - result = compare_imap_folders.get_email_count(mock_conn, "INBOX") - assert result is None - - def test_imap_exception_in_count(self): - """Test exception handling in get_email_count.""" - mock_conn = MagicMock() - mock_conn.select.side_effect = imaplib.IMAP4.error("Crash") - - result = compare_imap_folders.get_email_count(mock_conn, "INBOX") - assert result is None - - class TestConfigValidation: """Tests for configuration validation.""" diff --git a/test/test_count_imap_emails.py b/test/test_count_imap_emails.py index ca46c73..e1d893f 100644 --- a/test/test_count_imap_emails.py +++ b/test/test_count_imap_emails.py @@ -9,16 +9,15 @@ Tests cover: - Configuration validation """ -import imaplib import os import sys -from unittest.mock import MagicMock import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) import count_imap_emails +import imap_common from conftest import make_single_mock_connection @@ -98,102 +97,154 @@ class TestLocalEmailCounting: assert "TOTAL" in captured.out assert "3" in captured.out + def test_count_local_ignores_hidden_dirs(self, tmp_path, capsys): + inbox_path = tmp_path / "INBOX" + inbox_path.mkdir() + (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody") + (inbox_path / "note.txt").write_text("ignore") -class TestEmailCountingErrors: - """Tests for error handling in email counting.""" + hidden_path = tmp_path / ".hidden" + hidden_path.mkdir() + (hidden_path / "1_hidden.eml").write_bytes(b"Subject: Hidden\r\n\r\nBody") - def test_connection_failure(self, monkeypatch): - """Test graceful exit when connection fails.""" - mock_get = MagicMock(return_value=None) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) + cache_path = tmp_path / "__pycache__" + cache_path.mkdir() + (cache_path / "1_cache.eml").write_bytes(b"Subject: Cache\r\n\r\nBody") - # Should return silently without raising - count_imap_emails.count_emails("host", "user", "pass") - mock_get.assert_called_once() + nested_path = tmp_path / "Projects" / "Sub" + nested_path.mkdir(parents=True) + (nested_path / "1_sub.eml").write_bytes(b"Subject: Sub\r\n\r\nBody") - def test_list_command_failure(self, monkeypatch): - """Test handling of LIST command failure.""" - mock_mail = MagicMock() - mock_mail.list.return_value = ("NO", []) - mock_get = MagicMock(return_value=mock_mail) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) - - count_imap_emails.count_emails("host", "user", "pass") - mock_mail.list.assert_called_once() - # Should exit early, so select should not be called - mock_mail.select.assert_not_called() - - def test_select_command_failure(self, monkeypatch): - """Test handling of SELECT command failure for a folder.""" - mock_mail = MagicMock() - mock_mail.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"']) - # Fail selection - mock_mail.select.return_value = ("NO", [b"Select failed"]) - - mock_get = MagicMock(return_value=mock_mail) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) - monkeypatch.setattr("imap_common.normalize_folder_name", lambda x: "INBOX") - - count_imap_emails.count_emails("host", "user", "pass") - - mock_mail.select.assert_called_once() - # Should skip search for this folder - mock_mail.search.assert_not_called() - - def test_search_command_failure(self, monkeypatch, capsys): - """Test handling of SEARCH command failure.""" - mock_mail = MagicMock() - mock_mail.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"']) - mock_mail.select.return_value = ("OK", [b"Selected"]) - # Fail search - mock_mail.search.return_value = ("NO", [b"Search failed"]) - - mock_get = MagicMock(return_value=mock_mail) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) - monkeypatch.setattr("imap_common.normalize_folder_name", lambda x: "INBOX") - - count_imap_emails.count_emails("host", "user", "pass") + count_imap_emails.count_local_emails(str(tmp_path)) captured = capsys.readouterr() - assert "Error" in captured.out + assert "INBOX" in captured.out + assert "Projects/Sub" in captured.out + assert ".hidden" not in captured.out + assert "__pycache__" not in captured.out - def test_imap_exception_during_list(self, monkeypatch, capsys): - """Test handling of IMAP4 exception during list command.""" - mock_mail = MagicMock() - mock_mail.list.side_effect = imaplib.IMAP4.error("Crash listing") - mock_get = MagicMock(return_value=mock_mail) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) + def test_get_local_email_count_unreadable_folder(self, tmp_path): + inbox_path = tmp_path / "INBOX" + inbox_path.mkdir() + (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody") - count_imap_emails.count_emails("host", "user", "pass") + os.chmod(inbox_path, 0) + try: + result = imap_common.get_local_email_count(str(tmp_path), "INBOX") + assert result is None + finally: + os.chmod(inbox_path, 0o700) - captured = capsys.readouterr() - assert "Failed to list mailboxes" in captured.out - def test_imap_exception_during_select(self, monkeypatch, capsys): - """Test handling of IMAP4 exception during folder selection.""" - mock_mail = MagicMock() - mock_mail.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"']) - mock_mail.select.side_effect = imaplib.IMAP4.error("Crash selecting") +class TestImapCommonHelpers: + """Tests for imap_common helpers via script tests.""" - mock_get = MagicMock(return_value=mock_mail) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) - monkeypatch.setattr("imap_common.normalize_folder_name", lambda x: "INBOX") + def test_list_selectable_folders_filters_noselect(self): + class FakeConn: + def list(self): + return ( + "OK", + [ + b'(\\Noselect) "/" "Archive"', + b'(\\HasNoChildren) "/" "INBOX"', + '(\\HasNoChildren) "/" "Sent"', + ], + ) - count_imap_emails.count_emails("host", "user", "pass") + result = imap_common.list_selectable_folders(FakeConn()) + assert result == ["INBOX", "Sent"] - captured = capsys.readouterr() - # Should print Error for that folder - assert "Error" in captured.out + def test_list_selectable_folders_list_error(self): + class FakeConn: + def list(self): + return ("NO", []) - def test_generic_exception(self, monkeypatch, capsys): - """Test handling of generic connection/runtime exceptions.""" - mock_get = MagicMock(side_effect=Exception("Generic Crash")) - monkeypatch.setattr("imap_common.get_imap_connection", mock_get) + result = imap_common.list_selectable_folders(FakeConn()) + assert result == [] - count_imap_emails.count_emails("host", "user", "pass") + def test_list_selectable_folders_exception(self): + class FakeConn: + def list(self): + raise Exception("list failed") - captured = capsys.readouterr() - assert "An error occurred: Generic Crash" in captured.out + result = imap_common.list_selectable_folders(FakeConn()) + assert result == [] + + def test_get_imap_connection_oauth2_uses_authenticate(self, monkeypatch): + class FakeIMAP: + def __init__(self, _host): + self.auth_called = False + self.login_called = False + + def authenticate(self, _mechanism, auth_cb): + self.auth_called = True + auth_cb(None) + + def login(self, _user, _password): + self.login_called = True + + monkeypatch.setattr(imap_common.imaplib, "IMAP4_SSL", FakeIMAP) + + conn = imap_common.get_imap_connection("host", "user", oauth2_token="token") + + assert isinstance(conn, FakeIMAP) + assert conn.auth_called is True + assert conn.login_called is False + + def test_get_imap_connection_basic_login(self, monkeypatch): + class FakeIMAP: + def __init__(self, _host): + self.auth_called = False + self.login_called = False + + def authenticate(self, _mechanism, _auth_cb): + self.auth_called = True + + def login(self, _user, _password): + self.login_called = True + + monkeypatch.setattr(imap_common.imaplib, "IMAP4_SSL", FakeIMAP) + + conn = imap_common.get_imap_connection("host", "user", password="pass") + + assert isinstance(conn, FakeIMAP) + assert conn.login_called is True + assert conn.auth_called is False + + def test_ensure_connection_returns_same_conn_when_healthy(self): + class GoodConn: + def __init__(self): + self.noop_calls = 0 + + def noop(self): + self.noop_calls += 1 + + conn = GoodConn() + result = imap_common.ensure_connection(conn, "host", "user", "pass") + assert result is conn + assert conn.noop_calls == 1 + + def test_ensure_connection_reconnects_on_noop_error(self, monkeypatch): + class BadConn: + def noop(self): + raise Exception("fail") + + new_conn = object() + monkeypatch.setattr(imap_common, "get_imap_connection", lambda *args, **kwargs: new_conn) + + result = imap_common.ensure_connection(BadConn(), "host", "user", "pass") + assert result is new_conn + + def test_ensure_connection_from_conf_reconnects_on_noop_error(self, monkeypatch): + class BadConn: + def noop(self): + raise Exception("fail") + + new_conn = object() + monkeypatch.setattr(imap_common, "get_imap_connection_from_conf", lambda _conf: new_conn) + + result = imap_common.ensure_connection_from_conf(BadConn(), {"host": "h", "user": "u"}) + assert result is new_conn class TestMainFunction: diff --git a/test/test_migrate_imap_emails.py b/test/test_migrate_imap_emails.py index 39294a6..47bdd6b 100644 --- a/test/test_migrate_imap_emails.py +++ b/test/test_migrate_imap_emails.py @@ -19,6 +19,7 @@ import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) +import imap_common import migrate_imap_emails from conftest import make_mock_connection @@ -356,6 +357,82 @@ class TestGmailModeLabels: assert "Work" in dest_server.folders assert len(dest_server.folders["Work"]) == 2 + def test_gmail_mode_fallback_folder_for_unlabeled(self, mock_server_factory, monkeypatch): + msg = b"Subject: Unlabeled\r\nMessage-ID: \r\n\r\nBody" + + src_data = { + "[Gmail]/All Mail": [msg], + } + dest_data = {"INBOX": []} + + _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + env = { + "SRC_IMAP_HOST": "localhost", + "SRC_IMAP_USERNAME": "src_user", + "SRC_IMAP_PASSWORD": "p", + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "dest_user", + "DEST_IMAP_PASSWORD": "p", + "MAX_WORKERS": "1", + "GMAIL_MODE": "true", + } + monkeypatch.setattr(os, "environ", env) + monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2)) + + migrate_imap_emails.main() + + assert imap_common.FOLDER_RESTORED_UNLABELED in dest_server.folders + assert len(dest_server.folders[imap_common.FOLDER_RESTORED_UNLABELED]) == 1 + + +class TestCacheHitWithoutLock: + """Covers cached skip when no lock is provided.""" + + def test_cached_skip_without_lock(self, mock_server_factory): + msg_id = "" + msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + dest_data = {"INBOX": []} + + src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + src = imaplib.IMAP4("localhost", p1) + dest = imaplib.IMAP4("localhost", p2) + src.login("src_user", "p") + dest.login("dest_user", "p") + + src.select('"INBOX"', readonly=False) + + existing_dest_msg_ids = {msg_id} + success, _src, _dest, deleted = migrate_imap_emails.process_single_uid( + src, + dest, + b"1", + "INBOX", + False, + None, + False, + False, + None, + True, + False, + existing_dest_msg_ids=existing_dest_msg_ids, + existing_dest_msg_ids_lock=None, + progress_cache_path=None, + progress_cache_data=None, + progress_cache_lock=None, + dest_host="localhost", + dest_user="dest_user", + ) + + assert success is True + assert deleted == 0 + assert len(dest_server.folders["INBOX"]) == 0 + + src.logout() + dest.logout() + class TestConfigValidation: """Tests for configuration validation.""" @@ -454,6 +531,32 @@ class TestMigrateErrorHandling: migrate_imap_emails.main() + def test_select_error_in_process_batch(self, mock_server_factory, monkeypatch): + """Cover process_batch select exception handling with real server data.""" + src_data = {"INBOX": [b"Subject: Test\r\nMessage-ID: <1>\r\n\r\nBody"]} + dest_data = {"INBOX": []} + + _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + def raise_select(_self, _mailbox, readonly=False): + raise RuntimeError("Select failed") + + monkeypatch.setattr(imaplib.IMAP4, "select", raise_select) + monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2)) + + src_conf = {"host": "localhost", "user": "src_user", "password": "p"} + dest_conf = {"host": "localhost", "user": "dest_user", "password": "p"} + + migrate_imap_emails.process_batch( + [b"1"], + "INBOX", + src_conf, + dest_conf, + delete_from_source=False, + preserve_flags=False, + gmail_mode=False, + ) + def test_fetch_error_in_worker(self, mock_server_factory, monkeypatch): """Test error handling when fetching message details fails.""" src_data = {"INBOX": [b"Subject: Test\r\nMessage-ID: <1>\r\n\r\nBody"]} @@ -506,6 +609,35 @@ class TestMigrateErrorHandling: migrate_imap_emails.main() assert exc.value.code == 1 + def test_main_logs_progress_cache_load_failure(self, mock_server_factory, monkeypatch, capsys): + """Cover progress cache load exception in main.""" + src_data = {"INBOX": [b"Subject: X\r\nMessage-ID: \r\n\r\nBody"]} + dest_data = {"INBOX": []} + + _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2)) + monkeypatch.setattr( + imap_common, "load_progress_cache", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")) + ) + + env = { + "SRC_IMAP_HOST": "localhost", + "SRC_IMAP_USERNAME": "src_user", + "SRC_IMAP_PASSWORD": "p", + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "dest_user", + "DEST_IMAP_PASSWORD": "p", + "MAX_WORKERS": "1", + } + monkeypatch.setattr(os, "environ", env) + monkeypatch.setattr(sys, "argv", ["migrate_imap_emails.py", "--migrate-cache", "./cache"]) + + migrate_imap_emails.main() + + captured = capsys.readouterr() + assert "Warning: Failed to load progress cache" in captured.out + class TestTrashHandling: """Tests for trash folder related logic.""" @@ -571,6 +703,58 @@ class TestTrashHandling: assert len(src_server.folders["Trash"]) == 1 +class TestCommonMessageParsing: + """Covers imap_common message parsing helpers used by migrate.""" + + def test_parse_message_id_from_empty_bytes(self): + assert imap_common.parse_message_id_from_bytes(b"") is None + + def test_parse_message_id_and_subject_from_empty_bytes(self): + msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(b"") + assert msg_id is None + assert subject == "(No Subject)" + + def test_get_uid_to_message_id_map_empty(self): + result = imap_common.get_uid_to_message_id_map(object(), []) + assert result == {} + + def test_extract_message_id_invalid_type(self): + assert imap_common.extract_message_id(123) is None + + def test_parse_message_id_from_invalid_type(self): + assert imap_common.parse_message_id_from_bytes(123) is None + + def test_parse_message_id_from_bytes_success(self): + raw_message = b"Subject: X\r\nMessage-ID: \r\n\r\nBody" + assert imap_common.parse_message_id_from_bytes(raw_message) == "" + + def test_parse_message_id_and_subject_from_invalid_type(self): + msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(123) + assert msg_id is None + assert subject == "(No Subject)" + + def test_get_uid_to_message_id_map_missing_uid(self): + class FakeConn: + def uid(self, _cmd, _uids, _opts): + return ( + "OK", + [ + ( + b"1 (BODY[HEADER.FIELDS (MESSAGE-ID)] {40}", + b"Message-ID: \r\n", + ), + b")", + ], + ) + + result = imap_common.get_uid_to_message_id_map(FakeConn(), [b"1"]) + assert result == {} + + def test_decode_mime_header_exception_path(self): + result = imap_common.decode_mime_header(["not", "a", "header"]) + assert result == "['not', 'a', 'header']" + + class TestFilterPreservableFlags: """Tests for filter_preservable_flags function.""" @@ -733,3 +917,41 @@ class TestDestDeleteFunctionality: # All dest emails should be deleted assert len(dest_server.folders["INBOX"]) == 0 + + def test_dest_delete_syncs_after_migration(self, mock_server_factory, monkeypatch): + """End-to-end: delete orphans after a successful migration batch.""" + src_data = {"INBOX": [b"Subject: Keep\r\nMessage-ID: \r\n\r\nBody"]} + dest_data = { + "INBOX": [ + b"Subject: Keep\r\nMessage-ID: \r\n\r\nBody", + b"Subject: Orphan\r\nMessage-ID: \r\n\r\nBody", + ] + } + + src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2)) + + src = imaplib.IMAP4("localhost", p1) + dest = imaplib.IMAP4("localhost", p2) + src.login("src_user", "p") + dest.login("dest_user", "p") + + migrate_imap_emails.MAX_WORKERS = 1 + migrate_imap_emails.BATCH_SIZE = 1 + + migrate_imap_emails.migrate_folder( + src, + dest, + "INBOX", + False, + {"host": "localhost", "user": "src_user", "password": "p"}, + {"host": "localhost", "user": "dest_user", "password": "p"}, + dest_delete=True, + ) + + assert len(dest_server.folders["INBOX"]) == 1 + assert b"Message-ID: " in dest_server.folders["INBOX"][0]["content"] + + src.logout() + dest.logout() diff --git a/test/test_migrate_with_cache.py b/test/test_migrate_with_cache.py new file mode 100644 index 0000000..9f97d21 --- /dev/null +++ b/test/test_migrate_with_cache.py @@ -0,0 +1,248 @@ +"""End-to-end tests for migrate_imap_emails.py cache behavior.""" + +import imaplib +import json +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) + +import imap_common +import migrate_imap_emails +import restore_cache +from conftest import make_mock_connection + + +def _run_migrate(monkeypatch, cache_dir, src_port, dest_port, full_migrate=False, extra_env=None): + env = { + "SRC_IMAP_HOST": "localhost", + "SRC_IMAP_USERNAME": "src", + "SRC_IMAP_PASSWORD": "p", + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "dest", + "DEST_IMAP_PASSWORD": "p", + "MAX_WORKERS": "1", + } + if extra_env: + env.update(extra_env) + monkeypatch.setattr(os, "environ", env) + + argv = [ + "migrate_imap_emails.py", + "--src-host", + "localhost", + "--src-user", + "src", + "--src-pass", + "p", + "--dest-host", + "localhost", + "--dest-user", + "dest", + "--dest-pass", + "p", + "--migrate-cache", + str(cache_dir), + "--workers", + "1", + ] + if full_migrate: + argv.append("--full-migrate") + + monkeypatch.setattr(sys, "argv", argv) + monkeypatch.setattr( + migrate_imap_emails.imap_common, + "get_imap_connection", + make_mock_connection(src_port, dest_port, "src", "dest"), + ) + + migrate_imap_emails.main() + + +@pytest.mark.usefixtures("mock_server_factory") +class TestMigrationCache: + """End-to-end tests for incremental migration using local cache.""" + + def test_migrate_skips_cached_items(self, mock_server_factory, monkeypatch, tmp_path): + msg_id = "" + msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + dest_data = {"INBOX": []} + + _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + + # First run populates cache and copies message. + _run_migrate(monkeypatch, cache_dir, p1, p2) + assert len(dest_server.folders["INBOX"]) == 1 + + # Second run against a fresh destination should skip based on cache. + _src_server2, dest_server2, p3, p4 = mock_server_factory(src_data, {"INBOX": []}) + _run_migrate(monkeypatch, cache_dir, p3, p4) + assert len(dest_server2.folders["INBOX"]) == 0 + + def test_migrate_writes_to_cache(self, mock_server_factory, monkeypatch, tmp_path): + msg_id = "" + msg = f"Subject: New\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + dest_data = {"INBOX": []} + + _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + + _run_migrate(monkeypatch, cache_dir, p1, p2) + + assert len(dest_server.folders["INBOX"]) == 1 + + cache_path = restore_cache.get_dest_index_cache_path(str(cache_dir), "localhost", "dest") + assert os.path.exists(cache_path) + + with open(cache_path, encoding="utf-8") as f: + cache_data = json.load(f) + + msg_ids = set(cache_data.get("folders", {}).get("INBOX", {}).get("message_ids", [])) + assert msg_id in msg_ids + + def test_full_migrate_ignores_cache(self, mock_server_factory, monkeypatch, tmp_path): + msg_id = "" + msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + + # Populate cache with an initial run. + _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, {"INBOX": []}) + _run_migrate(monkeypatch, cache_dir, p1, p2) + + # Fresh destination should still copy when --full-migrate is set. + _src_server2, dest_server2, p3, p4 = mock_server_factory(src_data, {"INBOX": []}) + _run_migrate(monkeypatch, cache_dir, p3, p4, full_migrate=True) + + assert len(dest_server2.folders["INBOX"]) == 1 + + def test_cached_skip_with_preserve_flags(self, mock_server_factory, monkeypatch, tmp_path): + msg_id = "" + msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + + # Populate cache with initial run. + _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, {"INBOX": []}) + _run_migrate(monkeypatch, cache_dir, p1, p2) + + # Preserve flags disables pre-filtering, so cache skip happens per message. + _src_server2, dest_server2, p3, p4 = mock_server_factory(src_data, {"INBOX": []}) + _run_migrate(monkeypatch, cache_dir, p3, p4, extra_env={"PRESERVE_FLAGS": "true"}) + + assert len(dest_server2.folders["INBOX"]) == 0 + + def test_load_progress_cache_warns_on_unusable_root(self, tmp_path): + cache_file = tmp_path / "cachefile" + cache_file.write_text("not a directory") + + messages = [] + _cache_path, _cache_data, _cache_lock = imap_common.load_progress_cache( + str(cache_file), + "host", + "user", + log_fn=messages.append, + ) + + assert any("unable to create cache directory" in msg for msg in messages) + + def test_migrate_folder_logs_cache_load_failure(self, mock_server_factory, monkeypatch, tmp_path, capsys): + msg_id = "" + msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + dest_data = {"INBOX": []} + + _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2)) + monkeypatch.setattr( + imap_common, "load_progress_cache", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")) + ) + + src = imaplib.IMAP4("localhost", p1) + dest = imaplib.IMAP4("localhost", p2) + src.login("src", "p") + dest.login("dest", "p") + + monkeypatch.setattr(migrate_imap_emails, "get_thread_connections", lambda _src_conf, _dest_conf: (src, dest)) + + migrate_imap_emails.MAX_WORKERS = 1 + migrate_imap_emails.BATCH_SIZE = 1 + + migrate_imap_emails.migrate_folder( + src, + dest, + "INBOX", + False, + {"host": "localhost", "user": "src", "password": "p"}, + {"host": "localhost", "user": "dest", "password": "p"}, + progress_cache_path=str(tmp_path / "cache"), + progress_cache_file=None, + progress_cache_data=None, + progress_cache_lock=None, + ) + + captured = capsys.readouterr() + assert "Warning: Failed to load cache" in captured.out + assert len(dest_server.folders["INBOX"]) == 1 + + src.logout() + dest.logout() + + def test_migrate_folder_logs_cache_read_failure(self, mock_server_factory, monkeypatch, tmp_path, capsys): + msg_id = "" + msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode() + src_data = {"INBOX": [msg]} + dest_data = {"INBOX": []} + + _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data) + + monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2)) + monkeypatch.setattr( + restore_cache, + "get_cached_message_ids", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("read fail")), + ) + + src = imaplib.IMAP4("localhost", p1) + dest = imaplib.IMAP4("localhost", p2) + src.login("src", "p") + dest.login("dest", "p") + + monkeypatch.setattr(migrate_imap_emails, "get_thread_connections", lambda _src_conf, _dest_conf: (src, dest)) + + migrate_imap_emails.MAX_WORKERS = 1 + migrate_imap_emails.BATCH_SIZE = 1 + + migrate_imap_emails.migrate_folder( + src, + dest, + "INBOX", + False, + {"host": "localhost", "user": "src", "password": "p"}, + {"host": "localhost", "user": "dest", "password": "p"}, + progress_cache_path=str(tmp_path / "cache"), + progress_cache_file=None, + progress_cache_data=None, + progress_cache_lock=None, + ) + + captured = capsys.readouterr() + assert "Warning: Failed to read cache for folder 'INBOX'" in captured.out + assert len(dest_server.folders["INBOX"]) == 1 + + src.logout() + dest.logout() diff --git a/test/test_restore_imap_emails.py b/test/test_restore_imap_emails.py index c477316..eb43bff 100644 --- a/test/test_restore_imap_emails.py +++ b/test/test_restore_imap_emails.py @@ -9,10 +9,10 @@ Tests cover: - Configuration validation """ +import imaplib import json import os import sys -from unittest.mock import MagicMock import pytest @@ -147,6 +147,15 @@ Body content. assert message_id is None assert raw_content is None + def test_parse_eml_file_unknown_charset_subject(self, tmp_path): + """Test parsing with a subject that uses an unknown charset.""" + file_path = tmp_path / "unknown_charset.eml" + file_path.write_text("Subject: =?X-UNKNOWN?B?SGVsbG8=?=\r\nMessage-ID: \r\n\r\nBody") + + message_id, _date_str, _raw_content, subject = restore_imap_emails.parse_eml_file(str(file_path)) + assert message_id == "" + assert "Hello" in subject + class TestGetEmlFiles: """Tests for getting .eml files from a folder.""" @@ -175,46 +184,6 @@ class TestGetEmlFiles: assert result == [] -class TestGetBackupFolders: - """Tests for scanning backup folder structure.""" - - def test_get_backup_folders(self, tmp_path): - """Test scanning backup folders.""" - # Create folder structure - inbox = tmp_path / "INBOX" - inbox.mkdir() - (inbox / "email1.eml").write_text("content") - - sent = tmp_path / "Sent" - sent.mkdir() - (sent / "email2.eml").write_text("content") - - result = restore_imap_emails.get_backup_folders(str(tmp_path)) - - assert len(result) == 2 - folder_names = [f[0] for f in result] - assert "INBOX" in folder_names - assert "Sent" in folder_names - - def test_get_backup_folders_nested(self, tmp_path): - """Test scanning nested folder structure.""" - gmail = tmp_path / "[Gmail]" - gmail.mkdir() - all_mail = gmail / "All Mail" - all_mail.mkdir() - (all_mail / "email.eml").write_text("content") - - result = restore_imap_emails.get_backup_folders(str(tmp_path)) - - assert len(result) == 1 - assert result[0][0] == "[Gmail]/All Mail" - - def test_get_backup_folders_empty(self, tmp_path): - """Test scanning empty backup folder.""" - result = restore_imap_emails.get_backup_folders(str(tmp_path)) - assert result == [] - - class TestConfigValidation: """Tests for configuration validation.""" @@ -264,94 +233,6 @@ class TestConfigValidation: assert exc_info.value.code == 1 -class TestUploadEmail: - """Tests for email upload functionality.""" - - def test_upload_email_success(self, monkeypatch): - """Test successful email upload.""" - mock_conn = MagicMock() - mock_conn.create.return_value = ("OK", []) - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.append.return_value = ("OK", []) - - # Mock email_exists_in_folder to return False (not a duplicate) - monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False) - - result = restore_imap_emails.upload_email( - mock_conn, - "INBOX", - b"raw email content", - '"15-Jan-2024 10:30:00 +0000"', - "", - ) - - assert result == restore_imap_emails.UploadResult.SUCCESS - mock_conn.append.assert_called_once() - - def test_upload_email_duplicate(self, monkeypatch): - """Test upload returns ALREADY_EXISTS when message exists.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"1"]) - - # Mock email_exists_in_folder to return True (is a duplicate) - monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: True) - - result = restore_imap_emails.upload_email( - mock_conn, - "INBOX", - b"raw email content", - '"15-Jan-2024 10:30:00 +0000"', - "", - check_duplicate=True, - ) - - assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS - mock_conn.append.assert_not_called() - - def test_upload_email_with_seen_flag(self, monkeypatch): - """Test upload with \\Seen flag for read emails.""" - mock_conn = MagicMock() - mock_conn.create.return_value = ("OK", []) - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.append.return_value = ("OK", []) - - # Mock email_exists_in_folder to return False (not a duplicate) - monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False) - - result = restore_imap_emails.upload_email( - mock_conn, - "INBOX", - b"raw email content", - '"15-Jan-2024 10:30:00 +0000"', - "", - flags="\\Seen", # Mark as read - ) - - assert result == restore_imap_emails.UploadResult.SUCCESS - # Check that append was called with the \\Seen flag - call_args = mock_conn.append.call_args - assert call_args[0][1] == "(\\Seen)" - - def test_upload_email_failure(self, monkeypatch): - """Test upload returns FAILURE when an exception occurs.""" - mock_conn = MagicMock() - mock_conn.select.side_effect = Exception("Connection error") - - # Mock email_exists_in_folder to return False (not a duplicate) - monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False) - - result = restore_imap_emails.upload_email( - mock_conn, - "INBOX", - b"raw email content", - '"15-Jan-2024 10:30:00 +0000"', - "", - ) - - assert result == restore_imap_emails.UploadResult.FAILURE - mock_conn.append.assert_not_called() - - class TestRestoreIntegration: """Integration tests for restore functionality.""" @@ -391,6 +272,41 @@ Body content. # Run restore restore_imap_emails.main() + def test_restore_all_folders_scans_backup(self, single_mock_server, monkeypatch, tmp_path): + """End-to-end: restore all folders from a backup tree.""" + inbox = tmp_path / "INBOX" + inbox.mkdir() + (inbox / "1_Test_Email.eml").write_text("Subject: Inbox\nMessage-ID: \n\nBody") + + archive = tmp_path / "Archive" + archive.mkdir() + subfolder = archive / "Sub" + subfolder.mkdir() + (subfolder / "2_Test_Email.eml").write_text("Subject: Archive\nMessage-ID: \n\nBody") + + empty_folder = tmp_path / "Empty" + empty_folder.mkdir() + + dest_data = {"INBOX": []} + server, port = single_mock_server(dest_data) + + env = { + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "user", + "DEST_IMAP_PASSWORD": "pass", + } + monkeypatch.setattr(os, "environ", env) + monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "--src-path", str(tmp_path)]) + monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port)) + + restore_imap_emails.main() + + assert "INBOX" in server.folders + assert "Archive/Sub" in server.folders + assert len(server.folders["INBOX"]) == 1 + assert len(server.folders["Archive/Sub"]) == 1 + assert "Empty" not in server.folders + def test_restore_single_folder_full_restore_flag(self, single_mock_server, monkeypatch, tmp_path): """Smoke test: --full-restore flag is accepted.""" inbox = tmp_path / "INBOX" @@ -439,44 +355,70 @@ Body content. assert len(result) == 2 assert result[""] == ["INBOX", "Work"] + def test_restore_gmail_mode_fallback_folder(self, single_mock_server, monkeypatch, tmp_path): + """End-to-end: Gmail mode with no labels uses fallback folder.""" + gmail_all_mail = tmp_path / "[Gmail]" / "All Mail" + gmail_all_mail.mkdir(parents=True) + (gmail_all_mail / "1_Test.eml").write_text("Subject: X\nMessage-ID: \n\nBody") -class TestEmailExistsInFolder: - """Tests for duplicate detection.""" - - def test_email_exists_true(self, monkeypatch): - """Test detecting existing email.""" - mock_conn = MagicMock() - monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: True) - - result = restore_imap_emails.email_exists_in_folder(mock_conn, "") - assert result is True - - def test_email_exists_false(self, monkeypatch): - """Test detecting non-existing email.""" - mock_conn = MagicMock() - monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False) - - result = restore_imap_emails.email_exists_in_folder(mock_conn, "") - assert result is False - - def test_email_exists_no_message_id(self, monkeypatch): - """Test with no message ID.""" - mock_conn = MagicMock() - - result = restore_imap_emails.email_exists_in_folder(mock_conn, None) - assert result is False - - def test_email_exists_exception(self, monkeypatch): - """Test handling exception.""" - mock_conn = MagicMock() + dest_data = {"INBOX": []} + server, port = single_mock_server(dest_data) + env = { + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "user", + "DEST_IMAP_PASSWORD": "pass", + } + monkeypatch.setattr(os, "environ", env) monkeypatch.setattr( - "imap_common.message_exists_in_folder", - lambda *args: (_ for _ in ()).throw(Exception("Error")), + sys, + "argv", + ["restore_imap_emails.py", "--src-path", str(tmp_path), "--gmail-mode"], ) + monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port)) - result = restore_imap_emails.email_exists_in_folder(mock_conn, "") - assert result is False + restore_imap_emails.main() + + assert imap_common.FOLDER_RESTORED_UNLABELED in server.folders + assert len(server.folders[imap_common.FOLDER_RESTORED_UNLABELED]) == 1 + + def test_restore_dest_delete_removes_orphans(self, single_mock_server, monkeypatch, tmp_path): + """End-to-end: --dest-delete removes messages not in local backup.""" + inbox = tmp_path / "INBOX" + inbox.mkdir() + (inbox / "1_keep.eml").write_text("Subject: Keep\nMessage-ID: \n\nBody") + + dest_data = { + "INBOX": [ + b"Subject: Keep\r\nMessage-ID: \r\n\r\nBody", + b"Subject: Orphan\r\nMessage-ID: \r\n\r\nBody", + ] + } + server, port = single_mock_server(dest_data) + + env = { + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "user", + "DEST_IMAP_PASSWORD": "pass", + } + monkeypatch.setattr(os, "environ", env) + monkeypatch.setattr( + sys, + "argv", + [ + "restore_imap_emails.py", + "--src-path", + str(tmp_path), + "--dest-delete", + "INBOX", + ], + ) + monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port)) + + restore_imap_emails.main() + + assert len(server.folders["INBOX"]) == 1 + assert b"Message-ID: " in server.folders["INBOX"][0]["content"] class TestRestoreProgressCache: @@ -519,6 +461,60 @@ class TestRestoreProgressCache: assert "" in ids +class TestBackupFolderDiscovery: + """Tests for backup folder discovery helpers.""" + + def test_get_backup_folders_skips_unreadable(self, tmp_path): + inbox_path = tmp_path / "INBOX" + inbox_path.mkdir() + (inbox_path / "1.eml").write_bytes(b"Subject: Inbox\r\n\r\nBody") + + parent_path = tmp_path / "Parent" + parent_path.mkdir() + + unreadable_path = parent_path / "Unreadable" + unreadable_path.mkdir() + (unreadable_path / "1.eml").write_bytes(b"Subject: Hidden\r\n\r\nBody") + os.chmod(unreadable_path, 0) + + try: + folders = imap_common.get_backup_folders(str(tmp_path)) + finally: + os.chmod(unreadable_path, 0o700) + + folder_names = {name for name, _path in folders} + assert "INBOX" in folder_names + assert "Unreadable" not in folder_names + + def test_extract_message_id_from_eml_missing_file(self, tmp_path): + missing_path = tmp_path / "missing.eml" + assert imap_common.extract_message_id_from_eml(str(missing_path)) is None + + def test_extract_message_id_from_eml_success(self, tmp_path): + eml_path = tmp_path / "message.eml" + eml_path.write_text("Message-ID: \r\nSubject: Hi\r\n\r\nBody") + + assert imap_common.extract_message_id_from_eml(str(eml_path)) == "" + + +class TestTrashFolderDetection: + """Tests for trash folder detection with string LIST entries.""" + + def test_detect_trash_folder_with_string_entries(self): + class FakeConn: + def list(self): + return ( + "OK", + [ + '(\\HasNoChildren) "/" "INBOX"', + '(\\HasNoChildren \\Trash) "/" "Trash"', + ], + ) + + result = imap_common.detect_trash_folder(FakeConn()) + assert result == "Trash" + + class TestGetLabelsFromManifest: """Tests for get_labels_from_manifest function.""" @@ -590,68 +586,81 @@ class TestLabelToFolder: assert result == "Projects/2024" -class TestGmailModeDraftsFallbackRegression: - def test_gmail_mode_no_labels_does_not_upload_to_drafts(self, monkeypatch): - """Regression: messages with no usable labels must not be uploaded to Gmail Drafts.""" +class TestRestoreE2EHelpers: + """End-to-end tests for restore helper functions using the mock IMAP server.""" - captured = {} + def test_upload_email_success(self, single_mock_server): + dest_data = {"INBOX": []} + server, port = single_mock_server(dest_data) - def fake_upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True): - captured["folder_name"] = folder_name - return restore_imap_emails.UploadResult.SUCCESS + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") - def fake_parse_eml_file(_path): - return ( - "", - '"01-Jan-2024 00:00:00 +0000"', - b"Subject: X\r\nMessage-ID: \r\n\r\nBody", - "X", - ) - - monkeypatch.setattr(restore_imap_emails, "get_thread_connection", lambda _conf: MagicMock()) - monkeypatch.setattr(restore_imap_emails, "upload_email", fake_upload_email) - monkeypatch.setattr(restore_imap_emails, "parse_eml_file", fake_parse_eml_file) - - restore_imap_emails.process_restore_batch( - eml_files=[("/does/not/matter.eml", "x.eml")], - folder_name="__GMAIL_MODE__", - dest_conf=("host", "user", "pass"), - manifest={}, - apply_labels=True, - apply_flags=False, + result = restore_imap_emails.upload_email( + conn, + "INBOX", + b"Subject: Upload\r\nMessage-ID: \r\n\r\nBody", + '"15-Jan-2024 10:30:00 +0000"', + "", ) - assert captured["folder_name"] == imap_common.FOLDER_RESTORED_UNLABELED - assert captured["folder_name"] != "[Gmail]/Drafts" + assert result == restore_imap_emails.UploadResult.SUCCESS + assert len(server.folders["INBOX"]) == 1 + conn.logout() + def test_upload_email_duplicate(self, single_mock_server): + dest_data = {"INBOX": [b"Subject: Dup\r\nMessage-ID: \r\n\r\nBody"]} + server, port = single_mock_server(dest_data) -class TestSyncFlagsOnExisting: - """Tests for sync_flags_on_existing function.""" + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") - def test_sync_flags_adds_missing(self): - """Test that missing flags are added to existing email.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.search.return_value = ("OK", [b"1"]) - mock_conn.fetch.return_value = ("OK", [(b"1 (FLAGS ())", b"")]) - mock_conn.store.return_value = ("OK", None) + result = restore_imap_emails.upload_email( + conn, + "INBOX", + b"Subject: Dup\r\nMessage-ID: \r\n\r\nBody", + '"15-Jan-2024 10:30:00 +0000"', + "", + check_duplicate=True, + ) - # Should not raise - restore_imap_emails.sync_flags_on_existing(mock_conn, "INBOX", "", "\\Seen \\Flagged", 1000) + assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS + assert len(server.folders["INBOX"]) == 1 + conn.logout() - # Verify store was called with flags - mock_conn.store.assert_called() + def test_email_exists_in_folder(self, single_mock_server): + dest_data = {"INBOX": [b"Subject: Exists\r\nMessage-ID: \r\n\r\nBody"]} + _server, port = single_mock_server(dest_data) - def test_sync_flags_no_message_found(self): - """Test when message is not found.""" - mock_conn = MagicMock() - mock_conn.select.return_value = ("OK", [b"1"]) - mock_conn.search.return_value = ("OK", [b""]) # No match + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") + conn.select('"INBOX"') - # Should not raise or call store - restore_imap_emails.sync_flags_on_existing(mock_conn, "INBOX", "", "\\Seen", 1000) + assert restore_imap_emails.email_exists_in_folder(conn, "") is True + assert restore_imap_emails.email_exists_in_folder(conn, "") is False + conn.logout() - mock_conn.store.assert_not_called() + def test_sync_flags_on_existing(self, single_mock_server): + dest_data = { + "INBOX": [{"uid": 1, "flags": set(), "content": b"Subject: Flag\r\nMessage-ID: \r\n\r\nBody"}] + } + _server, port = single_mock_server(dest_data) + + conn = imaplib.IMAP4("localhost", port) + conn.login("user", "pass") + + restore_imap_emails.sync_flags_on_existing(conn, "INBOX", "", "\\Seen \\Flagged", 1000) + + conn.select('"INBOX"') + resp, data = conn.search(None, 'HEADER Message-ID ""') + assert resp == "OK" + msg_num = data[0].split()[0] + resp, flag_data = conn.fetch(msg_num, "(FLAGS)") + assert resp == "OK" + flag_text = str(flag_data[0]) + assert "\\Seen" in flag_text + assert "\\Flagged" in flag_text + conn.logout() class TestDestDeleteRestoreArgument: diff --git a/tools/mock_imap_server.py b/tools/mock_imap_server.py index c5dc4c6..193b6ea 100644 --- a/tools/mock_imap_server.py +++ b/tools/mock_imap_server.py @@ -80,6 +80,42 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): self.current_folders[folder] = [] self.send_response(tag, "OK CREATE completed") + elif cmd == "SEARCH": + if not self.selected_folder: + self.send_response(tag, RESPONSE_SELECT_FIRST) + continue + + msgs = self.current_folders[self.selected_folder] + sub_args = args + + header_msg_id = None + try: + m = re.search(r'HEADER\s+Message-ID\s+"([^"]+)"', sub_args, re.IGNORECASE) + if m: + header_msg_id = m.group(1) + except Exception: + header_msg_id = None + + seq_nums = [] + for idx, m in enumerate(msgs, start=1): + if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]: + continue + if header_msg_id: + msg_text = m["content"].decode("utf-8", errors="ignore") + if header_msg_id not in msg_text: + continue + seq_nums.append(str(idx)) + + if "ALL" in sub_args.upper() and not header_msg_id: + seq_nums = [str(idx) for idx in range(1, len(msgs) + 1)] + + seq_str = " ".join(seq_nums) + if seq_str: + self.wfile.write(f"* SEARCH {seq_str}\r\n".encode()) + else: + self.wfile.write(b"* SEARCH\r\n") + self.send_response(tag, RESPONSE_SEARCH_COMPLETED) + elif cmd == "EXPUNGE": if self.selected_folder: msgs = self.current_folders[self.selected_folder] @@ -125,7 +161,10 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): valid_uids.append(str(m["uid"])) uids_str = " ".join(valid_uids) - self.wfile.write(f"* SEARCH {uids_str}\r\n".encode()) + if uids_str: + self.wfile.write(f"* SEARCH {uids_str}\r\n".encode()) + else: + self.wfile.write(b"* SEARCH\r\n") self.send_response(tag, RESPONSE_SEARCH_COMPLETED) elif sub_cmd == "STORE": @@ -292,50 +331,6 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): print(f"MOCK APPEND ERROR: {e}") self.send_response(tag, "BAD APPEND") - elif cmd == "SEARCH": - # Parse SEARCH ALL, SEARCH HEADER Message-ID "...", etc. - # args might be: ALL, HEADER Message-ID "<123>", CHARSETS UTF-8 ... - - found_indices = [] - - if not self.selected_folder: - self.wfile.write(b"* SEARCH\r\n") - self.send_response(tag, RESPONSE_SEARCH_COMPLETED) - continue - - msgs = self.current_folders[self.selected_folder] - - if "ALL" in args.upper(): - # Return all message sequence numbers - found_indices = [str(idx + 1) for idx in range(len(msgs))] - elif "HEADER Message-ID" in args: - # Extract value - # Expected: ... HEADER Message-ID "value" ... - try: - # Split by 'MESSAGE-ID' (case insensitive?) - # part after Message-ID - post_mi = args.split("Message-ID", 1)[1].strip() - # Should start with quote or value - if post_mi.startswith('"'): - search_val = post_mi.split('"', 2)[1] - else: - search_val = post_mi.split(" ", 1)[0] - - search_val = search_val.replace("<", "").replace(">", "") - - for idx, m in enumerate(msgs): - content_str = m["content"].decode("utf-8", errors="ignore") - if search_val in content_str: - # Simple substring check is risky but okay for mock - # Better: regex for Message-ID: <...search_val...> - found_indices.append(str(idx + 1)) - except Exception: - pass - - indices_str = " ".join(found_indices) - self.wfile.write(f"* SEARCH {indices_str}\r\n".encode()) - self.send_response(tag, RESPONSE_SEARCH_COMPLETED) - elif cmd == "STORE": # STORE +FLAGS (\Seen) (non-UID; uses message sequence numbers) if not self.selected_folder: