Implement incremental restore with full restore option via restore cache. (#12)

* Implement incremental restore with full restore option and progress caching

* Refactor restore cache documentation and simplify cache path retrieval

* Add progress recording functionality for incremental restores

* Refactor email folder handling by introducing helper functions for folder creation and email appending

* Normalize flags format in append_email function for consistency

* Remove unnecessary blank line before get_labels_from_manifest function

* Enhance cache persistence with logging and error handling in save_dest_index_cache

* Add future annotations import for type hinting support

* Update src/restore_imap_emails.py

When applying Gmail labels, emails uploaded to label folders are not recorded in the progress cache. This means on subsequent incremental restore runs, the code will need to check the server again for these label folders even though the emails were just uploaded. Consider calling restore_cache.record_progress after successfully appending to a label folder (after line 514) to cache the Message-ID for that label folder as well.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/restore_cache.py

The condition for saving the cache will trigger a save based on the time threshold even when there are no pending updates (pending = 0). This can result in unnecessary disk writes. Consider adding a check to skip saving if there are no pending updates: force or (pending > 0 and (pending >= _MIN_PENDING_UPDATES_BEFORE_SAVE or (now - last_saved) >= _MIN_SECONDS_BETWEEN_SAVES)). This way, the cache is only saved when there are actually updates to persist, unless force=True.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/restore_imap_emails.py

'except' clause does nothing but pass and there is no explanatory comment.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Refactor condition in maybe_save_dest_index_cache for clarity

* [WIP] WIP Address feedback on incremental and full restore implementation (#13)

* Initial plan

* Optimize add_cached_message_id for O(1) lookups using set

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Cache set at entry level to avoid O(n) set construction on each add

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Move copy import to top of file per Python conventions

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Improve code style: condense comments and docstring

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Add batch Message-ID fetching for faster duplicate detection (#9)

* Add batch Message-ID fetching for faster duplicate detection

* Remove duplicate get_message_ids_in_folder in migrate_imap_emails

* Optimize orphan email deletion by reusing existing UID mappings

* Remove unused get_msg_details function and related test code

* Update src/restore_imap_emails.py

The call to restore_cache.record_progress is missing required keyword arguments. The function signature requires parameters like existing_dest_msg_ids, existing_dest_msg_ids_lock, progress_cache_path, progress_cache_data, progress_cache_lock, dest_host, and dest_user. This call should match the pattern used at line 456-467 where all required parameters are properly passed as keyword arguments.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/imap_common.py

The docstring states that the flags parameter is "Passed through to imaplib.IMAP4.append unchanged", but this is inaccurate. The function actually normalizes the flags by wrapping them in parentheses if they aren't already (lines 99-106). The docstring should be updated to reflect that the function normalizes flags to ensure they are in the correct IMAP format with parentheses.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Add unit tests for ensure_folder_exists and append_email helper functions (#14)

* Initial plan

* Add comprehensive unit tests for ensure_folder_exists and append_email

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Refactor folder existence checks to use imap_common.ensure_folder_exists directly

* Initial plan (#15)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

* Refactor process_restore_batch to improve argument clarity for restore_cache.record_progress

* Update src/imap_common.py

The append_email function catches all exceptions and returns False, but this makes it impossible for callers to distinguish between legitimate duplicate detection (email already exists) and actual errors (network failure, server error, etc.). Consider either raising exceptions for actual errors while returning False only for duplicates, or returning a more informative result (e.g., a tuple with status and reason, or an enum). This would allow restore_imap_emails.py to only record progress when emails are confirmed to be on the server.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Refactor test_append_exception to propagate exceptions for better error handling

* Update src/imap_common.py

The append_email function does not handle exceptions from imap_conn.append(). If the append call fails with an exception (e.g., connection error, permission error), the exception will propagate to the caller instead of returning False. This is inconsistent with the documented behavior that callers expect a boolean return value to indicate success or failure. The function should wrap the append call in a try-except block and return False on exception.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* [WIP] Address feedback on incremental restore implementation (#16)

* Initial plan

* Fix: Check append_email return value before recording progress for label application

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Add test to verify append_email return value is checked before recording progress

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Improve error message to clarify retry behavior for failed label applications

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* [WIP] Address feedback on incremental restore implementation (#17)

* Initial plan

* Fix progress recording for failed uploads with three-state result enum

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Clean up whitespace and simplify labels manifest definition in tests

* [WIP] Fix incremental restore return type consistency (#19)

* Initial plan

* Fix return type annotation for save_dest_index_cache to bool

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* [WIP] Work in progress on incremental restore feedback (#20)

* Initial plan

* Fix label folder progress tracking to use correct message ID set

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

* Refactor conditional statement for clarity in process_restore_batch function

* Remove exception handling in append_email function for cleaner flow

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>
Co-authored-by: Nathan Moinvaziri <nathan@nathanm.com>
This commit is contained in:
Javier Callico 2026-02-01 15:17:37 -05:00 committed by GitHub
parent 68c741913b
commit 7e853bdf27
8 changed files with 1047 additions and 55 deletions

View File

@ -543,6 +543,9 @@ This creates a `flags_manifest.json` with the status of each email.
### 9. Restore Backup to IMAP Server
Restore emails from a local backup to any IMAP server.
By default, restore runs in **incremental mode**: it uploads only emails that are not already present on the destination (based on `Message-ID`).
Use `--full-restore` to force the legacy behavior (process all emails and re-sync labels/flags for already-present messages).
```bash
# Restore all folders from backup
python3 restore_imap_emails.py \
@ -551,6 +554,14 @@ python3 restore_imap_emails.py \
--dest-user "you@gmail.com" \
--dest-pass "your-app-password"
# Force full restore (legacy behavior)
python3 restore_imap_emails.py \
--src-path "./my_backup" \
--dest-host "imap.gmail.com" \
--dest-user "you@gmail.com" \
--dest-pass "your-app-password" \
--full-restore
# Restore with flags (read/starred status)
python3 restore_imap_emails.py \
--src-path "./my_backup" \

View File

@ -4,6 +4,8 @@ IMAP Common Utilities
Shared functionality for IMAP migration, counting, and comparison scripts.
"""
from __future__ import annotations
import imaplib
import os
import re
@ -58,6 +60,55 @@ CMD_FETCH = "fetch"
OP_ADD_FLAGS = "+FLAGS"
def ensure_folder_exists(imap_conn, folder_name: str) -> None:
"""Best-effort create of a folder if it doesn't already exist.
IMAP servers typically return an error if the mailbox exists; this helper
intentionally ignores those errors.
"""
try:
if folder_name and folder_name.upper() != FOLDER_INBOX:
imap_conn.create(f'"{folder_name}"')
except Exception:
# Folder may already exist, or server may restrict creation.
pass
def append_email(
imap_conn,
folder_name: str,
raw_content: bytes,
date_str: str,
flags: str | None = None,
*,
ensure_folder: bool = True,
) -> bool:
"""Append an email message to a folder.
This is intentionally a thin wrapper around IMAP APPEND; callers can
perform duplicate checks or folder selection separately.
Args:
flags: Optional IMAP flags. If provided, they are normalized to a
parenthesized list before being passed to `imaplib.IMAP4.append`.
ensure_folder: If True, attempts to create the folder first (best-effort).
"""
if ensure_folder:
ensure_folder_exists(imap_conn, folder_name)
normalized_flags = None
if flags:
stripped = str(flags).strip()
if stripped:
if stripped.startswith("(") and stripped.endswith(")"):
normalized_flags = stripped
else:
normalized_flags = f"({stripped})"
resp, _ = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
return resp == "OK"
def verify_env_vars(vars_list):
"""
Checks if all environment variables in the list are set.

View File

@ -344,17 +344,10 @@ def process_batch(
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
return
def ensure_dest_folder(folder):
try:
if folder.upper() != imap_common.FOLDER_INBOX:
dest.create(f'"{folder}"')
except Exception:
pass
# In non-Gmail-mode, we keep a selected destination folder for efficiency
if not gmail_mode:
try:
ensure_dest_folder(folder_name)
imap_common.ensure_folder_exists(dest, folder_name)
dest.select(f'"{folder_name}"')
except Exception as e:
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
@ -424,7 +417,7 @@ def process_batch(
target_folder = folder_name
remaining_labels = []
ensure_dest_folder(target_folder)
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))
@ -435,7 +428,14 @@ def process_batch(
sync_flags_on_existing(dest, target_folder, msg_id, flags, size)
else:
valid_flags = f"({flags})" if (preserve_flags and flags) else None
dest.append(f'"{target_folder}"', valid_flags, date_str, msg_content)
imap_common.append_email(
dest,
target_folder,
msg_content,
date_str,
valid_flags,
ensure_folder=False,
)
safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}")
if preserve_flags and flags:
for flag in flags.split():
@ -450,11 +450,18 @@ def process_batch(
if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH):
continue
try:
ensure_dest_folder(label_folder)
imap_common.ensure_folder_exists(dest, label_folder)
dest.select(f'"{label_folder}"')
if not imap_common.message_exists_in_folder(dest, msg_id):
valid_flags = f"({flags})" if (preserve_flags and flags) else None
dest.append(f'"{label_folder}"', valid_flags, date_str, msg_content)
imap_common.append_email(
dest,
label_folder,
msg_content,
date_str,
valid_flags,
ensure_folder=False,
)
safe_print(f" -> Applied label: {label}")
if preserve_flags and flags:
for flag in flags.split():

254
src/restore_cache.py Normal file
View File

@ -0,0 +1,254 @@
"""Restore progress cache.
This module supports faster incremental restores by persisting, per destination+folder,
the set of Message-IDs already seen/processed by this tool.
The caller decides where the cache file lives by passing a cache directory/root.
"""
from __future__ import annotations
import copy
import json
import os
import re
import threading
import time
from collections.abc import Callable
RESTORE_CACHE_VERSION = 1
# Throttle disk writes so we can update frequently without rewriting a large JSON file
# on every single message.
_MIN_SECONDS_BETWEEN_SAVES = 2.0
_MIN_PENDING_UPDATES_BEFORE_SAVE = 50
def _safe_cache_component(value: str) -> str:
return re.sub(r"[^A-Za-z0-9_.-]+", "_", value or "")
def _prepare_cache_for_json(cache_data: dict) -> dict:
"""Create a deep copy of cache_data suitable for JSON serialization.
Removes in-memory helper fields like _ids_set that should not be persisted.
"""
snapshot = copy.deepcopy(cache_data)
folders = snapshot.get("folders", {})
if isinstance(folders, dict):
for folder_entry in folders.values():
if isinstance(folder_entry, dict):
# Remove the in-memory set cache
folder_entry.pop("_ids_set", None)
return snapshot
def get_dest_index_cache_path(cache_root: str, dest_host: str, dest_user: str) -> str:
safe_host = _safe_cache_component(dest_host)
safe_user = _safe_cache_component(dest_user)
return os.path.join(cache_root, f"restore_cache_{safe_host}_{safe_user}.json")
def load_dest_index_cache(cache_path: str) -> dict:
try:
with open(cache_path, encoding="utf-8") as f:
data = json.load(f)
if not isinstance(data, dict):
return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
if data.get("version") != RESTORE_CACHE_VERSION:
return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
if not isinstance(data.get("folders"), dict):
data["folders"] = {}
if not isinstance(data.get("_meta"), dict):
data["_meta"] = {}
return data
except FileNotFoundError:
return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
except Exception:
return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
def save_dest_index_cache(cache_path: str, cache_data: dict) -> bool:
try:
tmp_path = f"{cache_path}.tmp"
with open(tmp_path, "w", encoding="utf-8") as f:
json.dump(cache_data, f, ensure_ascii=False)
os.replace(tmp_path, cache_path)
return True
except Exception:
# Cache is best-effort; do not fail restore.
return False
def _ensure_dest(cache_data: dict, dest_host: str, dest_user: str) -> None:
cache_dest = cache_data.get("dest")
if not isinstance(cache_dest, dict) or cache_dest.get("host") != dest_host or cache_dest.get("user") != dest_user:
cache_data.clear()
cache_data.update(
{
"version": RESTORE_CACHE_VERSION,
"dest": {"host": dest_host, "user": dest_user},
"folders": {},
"_meta": {},
}
)
def get_cached_message_ids(
cache_data: dict,
cache_lock: threading.Lock,
dest_host: str,
dest_user: str,
folder_name: str,
) -> set[str]:
"""Return Message-IDs we've already seen/processed for this destination+folder."""
with cache_lock:
_ensure_dest(cache_data, dest_host, dest_user)
folders = cache_data.setdefault("folders", {})
entry = folders.get(folder_name)
if not isinstance(entry, dict):
return set()
ids = entry.get("message_ids")
if not isinstance(ids, list):
return set()
return {str(x) for x in ids if x}
def add_cached_message_id(
cache_data: dict,
cache_lock: threading.Lock,
dest_host: str,
dest_user: str,
folder_name: str,
message_id: str,
) -> bool:
"""Add a Message-ID to the cache. Returns True if it was newly added."""
if not message_id:
return False
msg_id = str(message_id).strip()
if not msg_id:
return False
with cache_lock:
_ensure_dest(cache_data, dest_host, dest_user)
folders = cache_data.setdefault("folders", {})
entry = folders.setdefault(folder_name, {})
if not isinstance(entry, dict):
folders[folder_name] = {}
entry = folders[folder_name]
ids = entry.get("message_ids")
if not isinstance(ids, list):
ids = []
entry["message_ids"] = ids
# Maintain entry-level in-memory set for O(1) lookups (not persisted to JSON)
ids_set = entry.get("_ids_set")
if ids_set is None:
ids_set = set(ids)
entry["_ids_set"] = ids_set
if msg_id in ids_set:
return False
ids.append(msg_id)
ids_set.add(msg_id)
meta = cache_data.setdefault("_meta", {})
if not isinstance(meta, dict):
meta = {}
cache_data["_meta"] = meta
meta["pending_updates"] = int(meta.get("pending_updates") or 0) + 1
return True
def maybe_save_dest_index_cache(
cache_path: str,
cache_data: dict,
cache_lock: threading.Lock,
*,
force: bool = False,
log_fn: Callable[[str], None] | None = None,
) -> bool:
"""Persist cache to disk if enough updates/time has accumulated."""
now = time.time()
pending = 0
with cache_lock:
meta = cache_data.setdefault("_meta", {})
if not isinstance(meta, dict):
meta = {}
cache_data["_meta"] = meta
pending = int(meta.get("pending_updates") or 0)
last_saved = float(meta.get("last_saved_ts") or 0.0)
should_save = force or (
pending > 0
and (pending >= _MIN_PENDING_UPDATES_BEFORE_SAVE or (now - last_saved) >= _MIN_SECONDS_BETWEEN_SAVES)
)
if not should_save:
return False
# Update meta before writing so the on-disk file reflects the flush decision.
meta["pending_updates"] = 0
meta["last_saved_ts"] = now
# Prepare a clean snapshot for JSON serialization (removes in-memory helper fields)
snapshot = _prepare_cache_for_json(cache_data)
did_write = save_dest_index_cache(cache_path, snapshot)
if did_write and log_fn is not None:
log_fn(f"Wrote restore cache ({pending} updates): {cache_path}")
return did_write
def record_progress(
*,
message_id: str | None,
folder_name: str,
existing_dest_msg_ids: set[str] | None,
existing_dest_msg_ids_lock: threading.Lock | None,
progress_cache_path: str | None,
progress_cache_data: dict | None,
progress_cache_lock: threading.Lock | None,
dest_host: str | None,
dest_user: str | None,
log_fn: Callable[[str], None] | None = None,
) -> None:
"""Record a processed Message-ID for fast skipping on future incremental runs.
Updates both:
- the in-memory set used by the current run, and
- the persisted progress cache on disk (throttled writes).
"""
if not message_id:
return
msg_id = str(message_id).strip()
if not msg_id:
return
if existing_dest_msg_ids is not None and existing_dest_msg_ids_lock is not None:
with existing_dest_msg_ids_lock:
existing_dest_msg_ids.add(msg_id)
if (
progress_cache_path
and progress_cache_data is not None
and progress_cache_lock is not None
and dest_host
and dest_user
):
add_cached_message_id(
progress_cache_data,
progress_cache_lock,
dest_host,
dest_user,
folder_name,
msg_id,
)
maybe_save_dest_index_cache(
progress_cache_path,
progress_cache_data,
progress_cache_lock,
log_fn=log_fn,
)

View File

@ -56,9 +56,21 @@ import time
from email import policy
from email.parser import BytesParser
from email.utils import parsedate_to_datetime
from enum import Enum
from typing import Optional
import imap_common
import imap_oauth2
import restore_cache
class UploadResult(Enum):
"""Result of an email upload operation."""
SUCCESS = "success"
ALREADY_EXISTS = "already_exists"
FAILURE = "failure"
# Defaults
MAX_WORKERS = 4 # Lower default for restore to avoid rate limits
@ -293,7 +305,7 @@ def email_exists_in_folder(imap_conn, message_id):
def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True):
"""
Upload a single email to the destination folder.
Returns True on success, False if duplicate or on failure.
Returns UploadResult enum: SUCCESS, ALREADY_EXISTS, or FAILURE.
Args:
flags: Optional string of IMAP flags like "\\Seen" for read emails.
@ -301,26 +313,29 @@ def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=Non
"""
try:
# Ensure folder exists
if folder_name.upper() != imap_common.FOLDER_INBOX:
try:
dest.create(f'"{folder_name}"')
except Exception:
pass # Folder may already exist
imap_common.ensure_folder_exists(dest, folder_name)
# Select folder
dest.select(f'"{folder_name}"')
# Check for duplicates if requested
if check_duplicate and message_id and email_exists_in_folder(dest, message_id):
return False # Already exists
return UploadResult.ALREADY_EXISTS
# Upload with original date and flags
resp, _ = dest.append(f'"{folder_name}"', flags, date_str, raw_content)
return resp == "OK"
success = imap_common.append_email(
dest,
folder_name,
raw_content,
date_str,
flags,
ensure_folder=False,
)
return UploadResult.SUCCESS if success else UploadResult.FAILURE
except Exception as e:
safe_print(f"Error uploading to {folder_name}: {e}")
return False
return UploadResult.FAILURE
def get_labels_from_manifest(manifest, message_id):
@ -351,7 +366,22 @@ def label_to_folder(label):
return label
def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_labels, apply_flags):
def process_restore_batch(
eml_files,
folder_name,
dest_conf,
manifest,
apply_labels,
apply_flags,
full_restore: bool = False,
existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None,
existing_dest_msg_ids_lock: Optional[threading.Lock] = None,
progress_cache_data: Optional[dict] = None,
progress_cache_lock: Optional[threading.Lock] = None,
progress_cache_path: Optional[str] = None,
dest_host: Optional[str] = None,
dest_user: Optional[str] = None,
):
"""
Process a batch of .eml files for restoration.
@ -415,28 +445,78 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab
target_folder = folder_name
remaining_labels = labels
# Upload to target folder
# In Gmail mode, check duplicates per-email since target folders vary
# In non-Gmail mode, duplicates were pre-filtered so no check needed
uploaded = upload_email(
dest, target_folder, raw_content, date_str, message_id, flags, check_duplicate=gmail_mode
)
existing_dest_msg_ids: Optional[set[str]] = None
if existing_dest_msg_ids_by_folder is not None:
if existing_dest_msg_ids_lock is None:
existing_dest_msg_ids_lock = threading.Lock()
if not uploaded:
with existing_dest_msg_ids_lock:
existing_dest_msg_ids = existing_dest_msg_ids_by_folder.get(target_folder)
# 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:
built = restore_cache.get_cached_message_ids(
progress_cache_data,
progress_cache_lock,
dest_host,
dest_user,
target_folder,
)
with existing_dest_msg_ids_lock:
existing_dest_msg_ids_by_folder.setdefault(target_folder, built)
existing_dest_msg_ids = existing_dest_msg_ids_by_folder[target_folder]
# Incremental default: if we already know we processed it before, skip entirely.
if (
not full_restore
and message_id
and existing_dest_msg_ids is not None
and message_id in existing_dest_msg_ids
):
safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
continue
# Upload to target folder.
# Keep server-side duplicate checks enabled to avoid creating duplicates for emails
# that exist on the destination but are not in our local progress cache.
upload_result = upload_email(dest, target_folder, raw_content, date_str, message_id, flags)
# Only record progress when upload succeeds or email already exists.
# Failed uploads should not be marked as processed to allow retry on next run.
if upload_result in (UploadResult.SUCCESS, UploadResult.ALREADY_EXISTS):
restore_cache.record_progress(
message_id=message_id,
folder_name=target_folder,
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,
)
if upload_result == UploadResult.ALREADY_EXISTS:
safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}")
# Even if skipped, sync flags on existing email if requested
if apply_flags and flags and message_id:
# Full restore preserves legacy behavior: sync flags on existing email if requested
if full_restore and apply_flags and flags and message_id:
sync_flags_on_existing(dest, target_folder, message_id, flags, size)
else:
elif upload_result == UploadResult.SUCCESS:
safe_print(f"[{target_folder}] UPLOADED | {size_str:<8} | {display_subject}")
# Show applied flags in same style as labels
if flags:
for flag in flags.split():
safe_print(f" -> Applied flag: {flag}")
else: # FAILURE
safe_print(f"[{target_folder}] FAILED | {size_str:<8} | {display_subject}")
# Apply remaining Gmail labels (always, whether uploaded or skipped)
# This ensures labels are synced even for existing emails
if apply_labels and remaining_labels:
# Apply remaining Gmail labels:
# - Full restore: apply/sync labels even for existing emails
# - Incremental (default): apply labels only for newly uploaded emails
if apply_labels and remaining_labels and (upload_result == UploadResult.SUCCESS or full_restore):
for label in remaining_labels:
label_folder = label_to_folder(label)
@ -454,19 +534,63 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab
try:
# Ensure label folder exists
if label_folder.upper() != imap_common.FOLDER_INBOX:
try:
dest.create(f'"{label_folder}"')
except Exception:
pass
imap_common.ensure_folder_exists(dest, label_folder)
# Select and check for duplicate
dest.select(f'"{label_folder}"')
if not email_exists_in_folder(dest, message_id):
dest.append(f'"{label_folder}"', flags, date_str, raw_content)
safe_print(f" -> Applied label: {label}")
# If email exists in this label folder, sync flags
elif apply_flags and flags:
append_success = imap_common.append_email(
dest,
label_folder,
raw_content,
date_str,
flags,
ensure_folder=False,
)
if append_success:
# Look up the correct set for label_folder
label_folder_msg_ids: Optional[set[str]] = None
if existing_dest_msg_ids_by_folder is not None:
with existing_dest_msg_ids_lock:
label_folder_msg_ids = existing_dest_msg_ids_by_folder.get(label_folder)
# Lazy-load per-folder progress cache if not yet loaded
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
and dest_host
and dest_user
):
built = restore_cache.get_cached_message_ids(
progress_cache_data,
progress_cache_lock,
dest_host,
dest_user,
label_folder,
)
with existing_dest_msg_ids_lock:
existing_dest_msg_ids_by_folder.setdefault(label_folder, built)
label_folder_msg_ids = existing_dest_msg_ids_by_folder[label_folder]
restore_cache.record_progress(
message_id=message_id,
folder_name=label_folder,
existing_dest_msg_ids=label_folder_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,
)
safe_print(f" -> Applied label: {label}")
else:
safe_print(f" -> Failed to apply label {label} (will retry on next restore)")
# If email exists in this label folder, sync flags (full restore only)
elif full_restore and apply_flags and flags:
sync_flags_on_existing(dest, label_folder, message_id, flags, size)
except Exception as e:
safe_print(f" -> Error applying label {label}: {e}")
@ -558,7 +682,17 @@ def delete_orphan_emails_from_dest(imap_conn, folder_name, local_msg_ids):
return deleted_count
def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_labels, apply_flags, dest_delete=False):
def restore_folder(
folder_name,
local_folder_path,
dest_conf,
manifest,
apply_labels,
apply_flags,
dest_delete=False,
full_restore: bool = False,
cache_root: Optional[str] = None,
):
"""
Restore all emails from a local folder to the destination IMAP server.
"""
@ -577,6 +711,30 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
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}")
# 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()}
existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
try:
existing_dest_msg_ids_by_folder[folder_name] = restore_cache.get_cached_message_ids(
cache_data,
cache_lock,
dest_conf["host"],
dest_conf["user"],
folder_name,
)
safe_print(f"Cache has {len(existing_dest_msg_ids_by_folder[folder_name])} Message-IDs for this folder.")
except Exception as e:
# Fall back to an empty cache for this folder if reading cached Message-IDs fails.
safe_print(f"Warning: Failed to load cached Message-IDs for folder '{folder_name}': {e}")
existing_dest_msg_ids_by_folder[folder_name] = set()
# If dest_delete enabled, get local Message-IDs for comparison
local_msg_ids = None
if dest_delete:
@ -641,6 +799,14 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
manifest,
apply_labels,
apply_flags,
full_restore,
existing_dest_msg_ids_by_folder,
existing_dest_msg_ids_lock,
cache_data,
cache_lock,
cache_path,
dest_conf.get("host"),
dest_conf.get("user"),
)
)
@ -658,8 +824,11 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids)
dest.logout()
# Force-flush progress cache at end.
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):
def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore: bool = False):
"""
Special restoration mode for Gmail: Upload emails to their first label folder
and then apply additional labels from the manifest.
@ -690,6 +859,19 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
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}")
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."
)
existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {} # lazily loaded per folder
existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
futures = []
for batch in batches:
@ -705,6 +887,14 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
manifest,
True, # apply_labels
apply_flags, # apply_flags
full_restore,
existing_dest_msg_ids_by_folder,
existing_dest_msg_ids_lock,
progress_cache_data,
progress_cache_lock,
cache_path,
dest_conf.get("host"),
dest_conf.get("user"),
)
)
@ -722,6 +912,9 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
elapsed = time.time() - start_time
safe_print(f"\nRestore completed in {elapsed:.1f}s")
# Force-flush progress cache at end.
restore_cache.maybe_save_dest_index_cache(cache_path, progress_cache_data, progress_cache_lock, force=True)
def get_backup_folders(local_path):
"""
@ -853,6 +1046,14 @@ def main():
help="Gmail restore mode: Upload to INBOX and apply labels + flags from manifest",
)
env_full_restore = os.getenv("FULL_RESTORE", "false").lower() == "true"
parser.add_argument(
"--full-restore",
action="store_true",
default=env_full_restore,
help="Force full restore (legacy): process all emails and sync labels/flags for already-present messages.",
)
# Sync mode: delete from dest emails not in local backup
env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
parser.add_argument(
@ -946,6 +1147,9 @@ def main():
print("Apply Flags : Yes (read/starred/answered/draft)")
if args.dest_delete:
print("Dest Delete : Yes (remove orphans from destination)")
print(
f"Restore Mode : {'Full (all emails)' if args.full_restore else 'Incremental (new emails only, use --full-restore to restore all)'}"
)
print("-----------------------------\n")
try:
@ -958,14 +1162,24 @@ def main():
if args.gmail_mode:
# Special Gmail mode
restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags)
restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore=args.full_restore)
elif args.folder:
# Restore specific folder
folder_path = os.path.join(local_path, args.folder.replace("/", os.sep))
if not os.path.exists(folder_path):
print(f"Error: Folder not found: {folder_path}")
sys.exit(1)
restore_folder(args.folder, folder_path, dest_conf, manifest, apply_labels, apply_flags, args.dest_delete)
restore_folder(
args.folder,
folder_path,
dest_conf,
manifest,
apply_labels,
apply_flags,
args.dest_delete,
full_restore=args.full_restore,
cache_root=local_path,
)
else:
# Restore all folders
folders = get_backup_folders(local_path)
@ -986,6 +1200,8 @@ def main():
apply_labels,
apply_flags,
args.dest_delete,
full_restore=args.full_restore,
cache_root=local_path,
)
print("\nRestore completed successfully.")

View File

@ -350,6 +350,290 @@ class TestExtractMessageId:
assert imap_common.extract_message_id(b"") is None
class TestEnsureFolderExists:
"""Tests for ensure_folder_exists function."""
def test_creates_folder_successfully(self):
"""Test folder is created when it doesn't exist."""
mock_conn = Mock()
mock_conn.create.return_value = ("OK", [])
# Should not raise any exception
imap_common.ensure_folder_exists(mock_conn, "TestFolder")
mock_conn.create.assert_called_once_with('"TestFolder"')
def test_folder_already_exists(self):
"""Test exception is suppressed when folder already exists."""
mock_conn = Mock()
mock_conn.create.side_effect = Exception("Folder already exists")
# Should not raise any exception
imap_common.ensure_folder_exists(mock_conn, "ExistingFolder")
mock_conn.create.assert_called_once_with('"ExistingFolder"')
def test_inbox_folder_skipped(self):
"""Test INBOX folder is not created."""
mock_conn = Mock()
imap_common.ensure_folder_exists(mock_conn, "INBOX")
mock_conn.create.assert_not_called()
# Also test lowercase
imap_common.ensure_folder_exists(mock_conn, "inbox")
mock_conn.create.assert_not_called()
def test_empty_folder_name(self):
"""Test empty folder name is skipped."""
mock_conn = Mock()
imap_common.ensure_folder_exists(mock_conn, "")
mock_conn.create.assert_not_called()
def test_none_folder_name(self):
"""Test None folder name is skipped."""
mock_conn = Mock()
imap_common.ensure_folder_exists(mock_conn, None)
mock_conn.create.assert_not_called()
def test_server_restriction_error_suppressed(self):
"""Test server restriction errors are suppressed."""
mock_conn = Mock()
mock_conn.create.side_effect = Exception("Permission denied")
# Should not raise any exception
imap_common.ensure_folder_exists(mock_conn, "RestrictedFolder")
mock_conn.create.assert_called_once()
class TestAppendEmail:
"""Tests for append_email function."""
def test_successful_append(self):
"""Test successful email append returns True."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
ensure_folder=False,
)
assert result is True
mock_conn.append.assert_called_once_with(
'"TestFolder"',
None,
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
def test_append_with_flags_with_parentheses(self):
"""Test flag normalization when flags already have parentheses."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
flags="(\\Seen \\Flagged)",
ensure_folder=False,
)
assert result is True
# Flags should remain unchanged
mock_conn.append.assert_called_once_with(
'"TestFolder"',
"(\\Seen \\Flagged)",
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
def test_append_with_flags_without_parentheses(self):
"""Test flag normalization when flags lack parentheses."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
flags="\\Seen",
ensure_folder=False,
)
assert result is True
# Flags should be wrapped in parentheses
mock_conn.append.assert_called_once_with(
'"TestFolder"',
"(\\Seen)",
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
def test_append_with_none_flags(self):
"""Test append with None flags."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
flags=None,
ensure_folder=False,
)
assert result is True
mock_conn.append.assert_called_once_with(
'"TestFolder"',
None,
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
def test_append_with_empty_flags(self):
"""Test append with empty string flags."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
flags="",
ensure_folder=False,
)
assert result is True
# Empty flags should result in None
mock_conn.append.assert_called_once_with(
'"TestFolder"',
None,
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
def test_append_with_whitespace_only_flags(self):
"""Test append with whitespace-only flags."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
flags=" ",
ensure_folder=False,
)
assert result is True
# Whitespace-only flags should result in None
mock_conn.append.assert_called_once_with(
'"TestFolder"',
None,
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
def test_append_with_ensure_folder_enabled(self):
"""Test append with ensure_folder enabled."""
mock_conn = Mock()
mock_conn.create.return_value = ("OK", [])
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
ensure_folder=True,
)
assert result is True
mock_conn.create.assert_called_once_with('"TestFolder"')
mock_conn.append.assert_called_once()
def test_append_with_ensure_folder_disabled(self):
"""Test append with ensure_folder disabled."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
ensure_folder=False,
)
assert result is True
mock_conn.create.assert_not_called()
mock_conn.append.assert_called_once()
def test_append_failure_returns_false(self):
"""Test append returns False on failure."""
mock_conn = Mock()
mock_conn.append.return_value = ("NO", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
ensure_folder=False,
)
assert result is False
def test_append_exception_propagates(self):
"""Test append propagates exceptions so callers can handle/log."""
mock_conn = Mock()
mock_conn.append.side_effect = Exception("Connection error")
import pytest
with pytest.raises(Exception, match="Connection error"):
imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
ensure_folder=False,
)
def test_append_with_multiple_flags(self):
"""Test append with multiple flags."""
mock_conn = Mock()
mock_conn.append.return_value = ("OK", [])
result = imap_common.append_email(
mock_conn,
"TestFolder",
b"email content",
"01-Jan-2024 12:00:00 +0000",
flags="\\Seen \\Flagged \\Answered",
ensure_folder=False,
)
assert result is True
# Multiple flags without parentheses should be wrapped
mock_conn.append.assert_called_once_with(
'"TestFolder"',
"(\\Seen \\Flagged \\Answered)",
"01-Jan-2024 12:00:00 +0000",
b"email content",
)
class TestGetMessageIdsInFolder:
"""Tests for get_message_ids_in_folder function."""

View File

@ -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 restore_cache
import restore_imap_emails
from conftest import make_single_mock_connection
@ -284,11 +285,11 @@ class TestUploadEmail:
"<test@test.com>",
)
assert result is True
assert result == restore_imap_emails.UploadResult.SUCCESS
mock_conn.append.assert_called_once()
def test_upload_email_duplicate(self, monkeypatch):
"""Test upload returns False when message exists."""
"""Test upload returns ALREADY_EXISTS when message exists."""
mock_conn = MagicMock()
mock_conn.select.return_value = ("OK", [b"1"])
@ -304,7 +305,7 @@ class TestUploadEmail:
check_duplicate=True,
)
assert result is False
assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS
mock_conn.append.assert_not_called()
def test_upload_email_with_seen_flag(self, monkeypatch):
@ -326,10 +327,29 @@ class TestUploadEmail:
flags="\\Seen", # Mark as read
)
assert result is True
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"
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:
@ -371,6 +391,39 @@ Body content.
# Run restore
restore_imap_emails.main()
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"
inbox.mkdir()
eml_content = b"""From: sender@test.com
To: recipient@test.com
Subject: Test Email
Message-ID: <test123@test.com>
Date: Mon, 15 Jan 2024 10:30:00 +0000
Body content.
"""
(inbox / "1_Test_Email.eml").write_bytes(eml_content)
src_data = {"INBOX": []}
_server, port = single_mock_server(src_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), "--full-restore", "INBOX"],
)
monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
restore_imap_emails.main()
def test_restore_with_labels_manifest(self, tmp_path):
"""Test that labels manifest is loaded correctly."""
# Create manifest
@ -416,6 +469,7 @@ class TestEmailExistsInFolder:
def test_email_exists_exception(self, monkeypatch):
"""Test handling exception."""
mock_conn = MagicMock()
monkeypatch.setattr(
"imap_common.message_exists_in_folder",
lambda *args: (_ for _ in ()).throw(Exception("Error")),
@ -425,6 +479,46 @@ class TestEmailExistsInFolder:
assert result is False
class TestRestoreProgressCache:
def test_progress_cache_add_get_persist(self, tmp_path):
import threading
cache_path = restore_cache.get_dest_index_cache_path(str(tmp_path), "imap.example.com", "user@example.com")
cache_data = restore_cache.load_dest_index_cache(cache_path)
lock = threading.Lock()
assert (
restore_cache.get_cached_message_ids(cache_data, lock, "imap.example.com", "user@example.com", "INBOX")
== set()
)
assert (
restore_cache.add_cached_message_id(
cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<a@test>"
)
is True
)
assert (
restore_cache.add_cached_message_id(
cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<a@test>"
)
is False
)
assert (
restore_cache.add_cached_message_id(
cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<b@test>"
)
is True
)
restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, lock, force=True)
cache_data2 = restore_cache.load_dest_index_cache(cache_path)
ids = restore_cache.get_cached_message_ids(cache_data2, lock, "imap.example.com", "user@example.com", "INBOX")
assert "<a@test>" in ids
assert "<b@test>" in ids
class TestGetLabelsFromManifest:
"""Tests for get_labels_from_manifest function."""
@ -504,7 +598,7 @@ class TestGmailModeDraftsFallbackRegression:
def fake_upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True):
captured["folder_name"] = folder_name
return True
return restore_imap_emails.UploadResult.SUCCESS
def fake_parse_eml_file(_path):
return (
@ -739,3 +833,67 @@ class TestDestDeleteRestoreFunctionality:
restore_imap_emails.main()
assert len(server.folders["INBOX"]) == 0
class TestAppendEmailReturnValueChecking:
"""Tests that verify append_email return values are checked before recording progress."""
def test_record_progress_not_called_on_append_failure(self, tmp_path, monkeypatch):
"""Test that record_progress is not called when append_email fails during label application."""
from unittest.mock import Mock, patch
# Create test email file
backup_root = tmp_path / "backup"
inbox = backup_root / "INBOX"
inbox.mkdir(parents=True)
eml_content = b"""From: sender@test.com
To: recipient@test.com
Subject: Test Email
Message-ID: <test123@test.com>
Date: Mon, 15 Jan 2024 10:30:00 +0000
Body content.
"""
(inbox / "1_Test_Email.eml").write_bytes(eml_content)
# Create labels manifest with multiple labels
manifest = {"<test123@test.com>": {"labels": ["INBOX", "Work"], "flags": []}}
manifest_path = backup_root / "labels_manifest.json"
manifest_path.write_text(json.dumps(manifest))
# Mock IMAP connection that fails on the second append (for label)
mock_conn = Mock()
mock_conn.select.return_value = ("OK", [b"0"])
mock_conn.search.return_value = ("OK", [b""]) # No duplicates
# First append succeeds (INBOX), second append fails (Work label)
mock_conn.append.side_effect = [
("OK", []), # INBOX upload succeeds
("NO", []), # Work label append fails
]
# Track calls to record_progress
with patch("restore_cache.record_progress") as mock_record_progress:
# Set up environment
monkeypatch.setenv("DEST_IMAP_HOST", "localhost")
monkeypatch.setenv("DEST_IMAP_USERNAME", "user")
monkeypatch.setenv("DEST_IMAP_PASSWORD", "pass")
monkeypatch.setenv("MAX_WORKERS", "1")
# Mock the connection
def mock_get_connection(*args, **kwargs):
return mock_conn
monkeypatch.setattr("imap_common.get_imap_connection", mock_get_connection)
monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "--src-path", str(backup_root), "INBOX"])
# Run restore
restore_imap_emails.main()
# Verify record_progress was called only once (for successful INBOX upload)
# and NOT called for the failed Work label append
assert mock_record_progress.call_count == 1
# Verify it was called for INBOX only
call_args = mock_record_progress.call_args
assert call_args[1]["folder_name"] == "INBOX"

View File

@ -105,10 +105,21 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
continue
msgs = self.current_folders[self.selected_folder]
# Support UID range searches like: UID SEARCH UID 101:*
uid_min = None
try:
m = re.search(r"UID\s+(\d+)", sub_args, re.IGNORECASE)
if m:
uid_min = int(m.group(1))
except Exception:
uid_min = None
# Handle UNDELETED
# If args has UNDELETED, filter flags
valid_uids = []
for m in msgs:
if uid_min is not None and m["uid"] < uid_min:
continue
if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]:
continue
valid_uids.append(str(m["uid"]))