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

@@ -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 <path>` 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.

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):

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)
 

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

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.

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()

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()

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: <msg1@test.com>\r\n"),
-                    b")",
-                    (b"2 (FLAGS () BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg2@test.com>\r\n"),
-                    b")",
-                    (
-                        b"3 (FLAGS (\\Seen \\Answered) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}",
-                        b"Message-ID: <msg3@test.com>\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: <msg1@test.com>\r\n\r\nBody",
+                b"Subject: B\r\nMessage-ID: <msg2@test.com>\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 "<msg1@test.com>" in result
         assert "<msg2@test.com>" in result
-        assert "<msg3@test.com>" 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: <msg1@test.com>\r\n"),
-                    b")",
-                    (b"2 (FLAGS () BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg2@test.com>\r\n"),
-                    b")",
-                    (
-                        b"3 (FLAGS (\\Seen \\Answered) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}",
-                        b"Message-ID: <msg3@test.com>\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: <msg1@test.com>\r\n"},
+                {"uid": 2, "flags": set(), "content": b"Message-ID: <msg2@test.com>\r\n"},
+                {"uid": 3, "flags": {"\\Seen", "\\Answered"}, "content": b"Message-ID: <msg3@test.com>\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 "<msg1@test.com>" in result
-        assert "\\Seen" in result["<msg1@test.com>"]["flags"]  # Has \Seen flag
+        assert "\\Seen" in result["<msg1@test.com>"]["flags"]
         assert "<msg2@test.com>" in result
-        assert result["<msg2@test.com>"]["flags"] == []  # No flags
+        assert result["<msg2@test.com>"]["flags"] == []
         assert "<msg3@test.com>" in result
-        assert "\\Seen" in result["<msg3@test.com>"]["flags"]  # Has \Seen flag
-        assert "\\Answered" in result["<msg3@test.com>"]["flags"]  # Also has \Answered
+        assert "\\Seen" in result["<msg3@test.com>"]["flags"]
+        assert "\\Answered" in result["<msg3@test.com>"]["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: <msg1@test.com>\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: <msg1@test.com>\r\n"],
+            "Work": [b"Message-ID: <msg1@test.com>\r\n"],
+            "[Gmail]/Sent Mail": [b"Message-ID: <msg2@test.com>\r\n"],
+            "[Gmail]/All Mail": [
+                {"uid": 1, "flags": {"\\Seen", "\\Flagged"}, "content": b"Message-ID: <msg1@test.com>\r\n"},
+                {"uid": 2, "flags": set(), "content": b"Message-ID: <msg2@test.com>\r\n"},
             ],
-        )
-
-        # Track which folder is selected
-        folder_data = {
-            "INBOX": {"<msg1@test.com>", "<msg2@test.com>"},
-            "Work": {"<msg1@test.com>"},
-            "[Gmail]/Sent Mail": {"<msg2@test.com>"},
         }
+        _server, port = single_mock_server(src_data)
 
-        # Mock info for All Mail (with flags)
-        all_mail_info = {
-            "<msg1@test.com>": {"flags": ["\\Seen", "\\Flagged"]},
-            "<msg2@test.com>": {"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 "<msg1@test.com>" in result
         assert "<msg2@test.com>" in result
         assert "INBOX" in result["<msg1@test.com>"]["labels"]
         assert "Work" in result["<msg1@test.com>"]["labels"]
         assert "\\Seen" in result["<msg1@test.com>"]["flags"]
         assert "\\Flagged" in result["<msg1@test.com>"]["flags"]
-        assert "INBOX" in result["<msg2@test.com>"]["labels"]
         assert "Sent Mail" in result["<msg2@test.com>"]["labels"]
         assert result["<msg2@test.com>"]["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."""

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."""
 

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:

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: <gm-unlabeled@test>\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 = "<cache-no-lock@test>"
+        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: <x@test>\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: <ok@test>\r\n\r\nBody"
+        assert imap_common.parse_message_id_from_bytes(raw_message) == "<ok@test>"
+
+    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: <x@test>\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: <keep@test>\r\n\r\nBody"]}
+        dest_data = {
+            "INBOX": [
+                b"Subject: Keep\r\nMessage-ID: <keep@test>\r\n\r\nBody",
+                b"Subject: Orphan\r\nMessage-ID: <orphan@test>\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: <keep@test>" in dest_server.folders["INBOX"][0]["content"]
+
+        src.logout()
+        dest.logout()

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 = "<cached@test>"
+        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 = "<new@test>"
+        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 = "<cached@test>"
+        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 = "<cached-preserve@test>"
+        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 = "<cache-fail@test>"
+        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 = "<cache-read-fail@test>"
+        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()

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: <unknown@test>\r\n\r\nBody")
+
+        message_id, _date_str, _raw_content, subject = restore_imap_emails.parse_eml_file(str(file_path))
+        assert message_id == "<unknown@test>"
+        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"',
-            "<test@test.com>",
-        )
-
-        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"',
-            "<test@test.com>",
-            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"',
-            "<test@test.com>",
-            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"',
-            "<test@test.com>",
-        )
-
-        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: <inbox@test>\n\nBody")
+
+        archive = tmp_path / "Archive"
+        archive.mkdir()
+        subfolder = archive / "Sub"
+        subfolder.mkdir()
+        (subfolder / "2_Test_Email.eml").write_text("Subject: Archive\nMessage-ID: <archive@test>\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["<msg1@test.com>"] == ["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: <no-labels@test>\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, "<test@test.com>")
-        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, "<test@test.com>")
-        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, "<test@test.com>")
-        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: <keep@test>\n\nBody")
+
+        dest_data = {
+            "INBOX": [
+                b"Subject: Keep\r\nMessage-ID: <keep@test>\r\n\r\nBody",
+                b"Subject: Orphan\r\nMessage-ID: <orphan@test>\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: <keep@test>" in server.folders["INBOX"][0]["content"]
 
 
 class TestRestoreProgressCache:
@@ -519,6 +461,60 @@ class TestRestoreProgressCache:
         assert "<b@test>" 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: <ok@test>\r\nSubject: Hi\r\n\r\nBody")
+
+        assert imap_common.extract_message_id_from_eml(str(eml_path)) == "<ok@test>"
+
+
+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 (
-                "<no-labels@test>",
-                '"01-Jan-2024 00:00:00 +0000"',
-                b"Subject: X\r\nMessage-ID: <no-labels@test>\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: <up@test>\r\n\r\nBody",
+            '"15-Jan-2024 10:30:00 +0000"',
+            "<up@test>",
         )
 
-        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: <dup@test>\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: <dup@test>\r\n\r\nBody",
+            '"15-Jan-2024 10:30:00 +0000"',
+            "<dup@test>",
+            check_duplicate=True,
+        )
 
-        # Should not raise
-        restore_imap_emails.sync_flags_on_existing(mock_conn, "INBOX", "<test@test.com>", "\\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: <exists@test>\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", "<test@test.com>", "\\Seen", 1000)
+        assert restore_imap_emails.email_exists_in_folder(conn, "<exists@test>") is True
+        assert restore_imap_emails.email_exists_in_folder(conn, "<missing@test>") 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: <flag@test>\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", "<flag@test>", "\\Seen \\Flagged", 1000)
+
+        conn.select('"INBOX"')
+        resp, data = conn.search(None, 'HEADER Message-ID "<flag@test>"')
+        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:

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 <msg_set> +FLAGS (\Seen)  (non-UID; uses message sequence numbers)
                     if not self.selected_folder: