gian/imap-migration-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add imap_session module for centralized OAuth2 refresh (#21)

Browse Source 
* Add .claude to gitignore

* Add imap_session module for centralized OAuth2 refresh

Create new imap_session.py module that combines imap_common and
imap_oauth2 for unified session management:

- ensure_connection(): Refreshes OAuth2 token and ensures connection health
- ensure_folder_session(): Ensures connection AND folder selection

Updated backup, migrate, and restore scripts to use the new module,
simplifying the proactive token refresh pattern in process_batch
functions.

* Move is_token_expired_error to imap_oauth2 module

Relocates the OAuth2 token expiration detection function to the oauth2-specific module where it logically belongs.

* Add retry logic for auth errors during email processing

Extract single-UID processing into dedicated functions with retry capability
for recoverable authentication failures. Broadens auth error detection beyond
token expiration to include session loss and authentication failures.

* Cache Microsoft tenant discovery results

Avoids repeated OpenID Connect discovery requests for the same email domain
during OAuth2 authentication flows. Normalizes domain to lowercase for
consistent cache keys.

* Add tests for is_token_expired_error and is_auth_error

Covers token expiration, session invalidation, and authentication failure
detection with case-insensitive matching verification. Ensures non-auth
errors like connection and timeout issues are correctly excluded.

* Add tests for imap_session module

Covers ensure_connection and ensure_folder_session functions including
OAuth2 token refresh, folder selection, and error handling scenarios.
Verifies connection change detection triggers folder re-selection.

* Add tests for auth error re-raising in get_message_ids_in_folder

Verifies that authentication errors are propagated to callers for
reconnection handling while non-auth errors return empty results.
Clarifies test names to distinguish error handling behaviors.
This commit is contained in:
Nathan Moinvaziri 2026-02-06 10:26:44 -08:00 committed by GitHub
parent 8c49bc7c48
commit e932190618
11 changed files with 887 additions and 236 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -7,3 +7,4 @@ env/
venv/
.env
.coverage
.claude

233
src/backup_imap_emails.py
View File

@ -58,6 +58,7 @@ import threading
import imap_common
import imap_oauth2
import imap_session
# Defaults
MAX_WORKERS = 10
@ -76,13 +77,6 @@ def safe_print(message):
print(f"[{short_name}] {message}")
def ensure_connection(conn, conf):
"""Refresh OAuth2 token if needed and ensure connection is healthy."""
if conf.get("oauth2"):
imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
return imap_common.ensure_connection_from_conf(conn, conf)
def get_thread_connection(src_conf):
# Backwards-compatible: tests and some internal call sites may still pass
# (host, user, password) tuples. Normalize to the dict form used by
@ -90,10 +84,74 @@ def get_thread_connection(src_conf):
if isinstance(src_conf, tuple) and len(src_conf) == 3:
src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]}
thread_local.src = ensure_connection(getattr(thread_local, "src", None), src_conf)
thread_local.src = imap_session.ensure_connection(getattr(thread_local, "src", None), src_conf)
return thread_local.src
def process_single_uid(src, uid, folder_name, local_folder_path):
"""
Fetch and save a single email by UID.
Args:
src: IMAP connection
uid: Message UID to fetch
folder_name: Current folder name (for logging)
local_folder_path: Local directory to save email
Returns:
Tuple of (success: bool, connection):
- (True, src): UID processed successfully (or skipped)
- (False, src): Auth error occurred, caller should retry after reconnect
"""
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
try:
resp, data = src.uid("fetch", uid, "(RFC822)")
if resp != "OK" or not data or data[0] is None:
safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
return (True, src) # Don't retry, move on
raw_email = None
for item in data:
if isinstance(item, tuple):
raw_email = item[1]
break
# Derive Subject for filename from the already-fetched message bytes.
_, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
if not subject:
clean_subject = "No Subject"
else:
clean_subject = imap_common.sanitize_filename(subject)
clean_subject = clean_subject[:100]
filename = f"{uid_str}_{clean_subject}.eml"
full_path = os.path.join(local_folder_path, filename)
if os.path.exists(full_path):
return (True, src) # Already exists
if raw_email:
try:
with open(full_path, "wb") as f:
f.write(raw_email)
safe_print(f"[{folder_name}] SAVED | {filename[:60]}...")
except OSError as e:
safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
else:
safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
return (True, src)
except Exception as e:
if imap_oauth2.is_auth_error(e):
safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
return (False, src) # Signal retry needed
else:
safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
return (True, src) # Don't retry other errors
def process_batch(uids, folder_name, src_conf, local_folder_path):
if isinstance(src_conf, tuple) and len(src_conf) == 3:
src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]}
@ -110,52 +168,24 @@ def process_batch(uids, folder_name, src_conf, local_folder_path):
return
for uid in uids:
try:
# Helper to handle byte UIDs
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
max_retries = 2
# Fetch Full Content
# RFC822 gets the whole message including headers and attachments
resp, data = src.uid("fetch", uid, "(RFC822)")
if resp != "OK" or not data or data[0] is None:
safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
continue
for attempt in range(max_retries):
src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True)
thread_local.src = src
if not ok:
safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}")
return
raw_email = None
for item in data:
if isinstance(item, tuple):
raw_email = item[1]
break
success, src = process_single_uid(src, uid, folder_name, local_folder_path)
thread_local.src = src
# Derive Subject for filename from the already-fetched message bytes.
_, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
if not subject:
clean_subject = "No Subject"
else:
clean_subject = imap_common.sanitize_filename(subject)
# Limit subject len to avoid FS path length limits (max 255 usually)
clean_subject = clean_subject[:100]
filename = f"{uid_str}_{clean_subject}.eml"
full_path = os.path.join(local_folder_path, filename)
# Double check existence (in case incremental UID checks missed it or naming collision).
if os.path.exists(full_path):
continue
if raw_email:
try:
with open(full_path, "wb") as f:
f.write(raw_email)
safe_print(f"[{folder_name}] SAVED | {filename[:60]}...")
except OSError as e:
# Valid filename might still fail if path too long
safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
else:
safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
except Exception as e:
safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
if success:
break
if attempt < max_retries - 1:
src = None
thread_local.src = None
def get_existing_uids(local_path):
@ -210,11 +240,21 @@ def is_gmail_label_folder(folder_name):
PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
"""
Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
Optional progress_callback(current, total) for progress reporting.
Returns a dict of Message-IDs and their IMAP flags for all emails in a folder,
with OAuth2 session management.
Args:
imap_conn: IMAP connection
folder_name: Folder to scan
src_conf: Connection config dict for OAuth2 refresh
progress_callback: Optional callback(current, total) for progress reporting
Returns:
Tuple of (message_info, imap_conn) where message_info is:
{ "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
The returned imap_conn may be different if reconnection occurred.
"""
message_info = {}
@ -222,16 +262,16 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
imap_conn.select(f'"{folder_name}"', readonly=True)
except Exception as e:
safe_print(f"Could not select folder {folder_name}: {e}")
return message_info
return (message_info, imap_conn)
try:
resp, data = imap_conn.uid("search", None, "ALL")
if resp != "OK" or not data or not data[0]:
return message_info
return (message_info, imap_conn)
uids = data[0].split()
if not uids:
return message_info
return (message_info, imap_conn)
total_uids = len(uids)
@ -245,6 +285,13 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
if progress_callback:
progress_callback(min(i + batch_size, total_uids), total_uids)
# Proactively refresh token and ensure folder is selected
if src_conf:
imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True)
if not ok:
safe_print(f"ERROR: Connection/folder lost in {folder_name}")
break
try:
# Fetch both Message-ID header and FLAGS
resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
@ -271,19 +318,44 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
message_info[msg_id] = {"flags": flags}
except Exception as e:
safe_print(f"Error fetching batch in {folder_name}: {e}")
# Try to keep connection alive
try:
imap_conn.noop()
except Exception:
pass
continue
except Exception as e:
safe_print(f"Error searching folder {folder_name}: {e}")
return (message_info, imap_conn)
def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
"""
Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
Optional progress_callback(current, total) for progress reporting.
"""
message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback)
return message_info
def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
"""
Returns a set of Message-IDs for all emails in a given folder,
with OAuth2 session management.
Args:
imap_conn: IMAP connection
folder_name: Folder to scan
src_conf: Connection config dict for OAuth2 refresh
progress_callback: Optional callback(current, total) for progress reporting
Returns:
Tuple of (message_ids, imap_conn) where message_ids is a set of strings.
The returned imap_conn may be different if reconnection occurred.
"""
info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback)
return (set(info.keys()), imap_conn)
def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
"""
Returns a set of Message-IDs for all emails in a given folder.
@ -294,7 +366,7 @@ def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
return set(info.keys())
def build_labels_manifest(imap_conn, local_path):
def build_labels_manifest(imap_conn, local_path, src_conf=None):
"""
Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags.
Scans all folders (labels) in the account and records which Message-IDs
@ -302,6 +374,7 @@ def build_labels_manifest(imap_conn, local_path):
Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
Saves the manifest to labels_manifest.json in the backup directory.
Optional src_conf for automatic token refresh on expiration.
"""
import time
@ -329,7 +402,9 @@ def build_labels_manifest(imap_conn, local_path):
flush=True,
)
all_mail_info = get_message_info_in_folder(imap_conn, imap_common.GMAIL_ALL_MAIL, all_mail_progress_cb)
all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
imap_conn, imap_common.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
)
print() # New line after progress
# Initialize manifest with flags from All Mail
@ -369,7 +444,7 @@ def build_labels_manifest(imap_conn, local_path):
)
safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
message_ids = get_message_ids_in_folder(imap_conn, folder_name, progress_cb)
message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
print() # New line after progress
# Keep connection alive between folders
@ -422,13 +497,14 @@ def build_labels_manifest(imap_conn, local_path):
return manifest
def build_flags_manifest(imap_conn, local_path, folders_to_scan=None):
def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None):
"""
Builds a manifest mapping Message-IDs to their IMAP flags.
For non-Gmail servers, scans specified folders (or all folders if not specified).
Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... }
Saves the manifest to flags_manifest.json in the backup directory.
Optional src_conf for automatic token refresh on expiration.
"""
import time
@ -467,7 +543,7 @@ def build_flags_manifest(imap_conn, local_path, folders_to_scan=None):
)
safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
folder_info = get_message_info_in_folder(imap_conn, folder_name, progress_cb)
folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
print() # New line after progress
# Merge into manifest (keep first occurrence of flags)
@ -829,7 +905,7 @@ def main():
if args.preserve_labels or args.manifest_only or args.gmail_mode:
print("Building Gmail labels manifest...")
print("This scans all folders to map Message-IDs to labels and flags.\n")
build_labels_manifest(src, local_path)
build_labels_manifest(src, local_path, src_conf)
print("") # Blank line after manifest building
# Build flags-only manifest for non-Gmail servers
elif args.preserve_flags:
@ -837,12 +913,15 @@ def main():
print("This scans folders to capture read/starred/etc status.\n")
# Get folders to scan
folders_to_scan = [args.folder] if args.folder else None
build_flags_manifest(src, local_path, folders_to_scan)
build_flags_manifest(src, local_path, folders_to_scan, src_conf)
print("") # Blank line after manifest building
# If manifest-only mode, we're done
if args.manifest_only:
src.logout()
try:
src.logout()
except Exception:
pass # Connection may already be closed
manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
print("\nManifest-only mode complete.")
print(f"Labels manifest saved to: {manifest_path}")
@ -856,15 +935,23 @@ def main():
elif args.folder:
backup_folder(src, args.folder, local_path, src_conf, args.dest_delete)
else:
# Reconnect after potentially long manifest building
src = imap_session.ensure_connection(src, src_conf)
if not src:
print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.")
sys.exit(0)
folders = imap_common.list_selectable_folders(src)
for name in folders:
src = ensure_connection(src, src_conf)
src = imap_session.ensure_connection(src, src_conf)
if not src:
print("Fatal: Could not reconnect to IMAP server. Aborting.")
sys.exit(1)
backup_folder(src, name, local_path, src_conf, args.dest_delete)
src.logout()
try:
src.logout()
except Exception:
pass # Connection may already be closed
print("\nBackup completed successfully.")
if args.preserve_labels or args.gmail_mode:

12
src/imap_common.py
View File

@ -14,6 +14,8 @@ from email import policy
from email.header import decode_header
from email.parser import BytesParser
import imap_oauth2
# Standard IMAP flags
FLAG_SEEN = "\\Seen"
FLAG_ANSWERED = "\\Answered"
@ -363,7 +365,10 @@ def get_message_ids_in_folder(imap_conn):
resp, data = imap_conn.uid("search", None, "ALL")
if resp != "OK" or not data[0].strip():
return {}
except Exception:
except Exception as e:
# Re-raise token expiration errors so callers can handle reconnection
if imap_oauth2.is_auth_error(e):
raise
return {}
uids = data[0].split()
@ -403,7 +408,10 @@ def get_uid_to_message_id_map(imap_conn, uids):
mid = extract_message_id(item[1])
if mid:
uid_to_msgid[uid] = mid
except Exception:
except Exception as e:
# Re-raise token expiration errors so callers can handle reconnection
if imap_oauth2.is_auth_error(e):
raise
continue
return uid_to_msgid

41
src/imap_oauth2.py
View File

@ -23,6 +23,7 @@ import urllib.parse
# Module-level caches for OAuth2 token refresh
_msal_app_cache = {} # (client_id, tenant_id) -> PublicClientApplication
_google_creds_cache = {} # (client_id, client_secret) -> credentials
_tenant_cache = {} # domain -> tenant_id
_token_refresh_lock = threading.Lock()
@ -46,6 +47,35 @@ def _fetch_json_https(host, path, timeout=10):
return json.loads(body.decode("utf-8"))
def is_token_expired_error(error):
"""
Check if an exception indicates OAuth2 token expiration.
Args:
error: The exception to check
Returns:
True if the error indicates token expiration, False otherwise
"""
error_str = str(error).lower()
return "accesstokenexpired" in error_str or "session invalidated" in error_str
def is_auth_error(error):
"""
Check if an exception indicates an authentication failure that may be recoverable
by reconnecting (token expiration, session loss, etc.).
Args:
error: The exception to check
Returns:
True if the error indicates an auth failure, False otherwise
"""
error_str = str(error).lower()
return is_token_expired_error(error) or "not authenticated" in error_str or "authentication failed" in error_str
def detect_oauth2_provider(host):
"""
Detects the OAuth2 provider from the IMAP host.
@ -63,13 +93,18 @@ def discover_microsoft_tenant(email):
"""
Auto-discovers the Microsoft tenant ID from an email address domain.
Uses the OpenID Connect discovery endpoint (no authentication required).
Results are cached per domain to avoid repeated network requests.
Returns the tenant ID string or None if discovery fails.
"""
domain = email.split("@")[-1].strip()
domain = email.split("@")[-1].strip().lower()
if not domain:
print("Error: Could not discover Microsoft tenant: missing email domain")
return None
# Return cached tenant if available
if domain in _tenant_cache:
return _tenant_cache[domain]
domain_quoted = urllib.parse.quote(domain, safe=".-")
path = f"/{domain_quoted}/.well-known/openid-configuration"
@ -82,7 +117,9 @@ def discover_microsoft_tenant(email):
issuer = data.get("issuer", "")
match = re.search(r"/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})", issuer)
if match:
return match.group(1)
tenant_id = match.group(1)
_tenant_cache[domain] = tenant_id
return tenant_id
print(f"Error: Could not extract tenant ID from issuer: {issuer}")
return None

62
src/imap_session.py Normal file
View File

@ -0,0 +1,62 @@
"""
IMAP Session Management
Connection and session management for IMAP operations with OAuth2 support.
Combines imap_common (low-level IMAP) with imap_oauth2 (token refresh).
"""
import imap_common
import imap_oauth2
def ensure_connection(conn, conf):
"""
Refresh OAuth2 token if needed and ensure connection is healthy.
Args:
conn: Existing IMAP connection or None
conf: Connection config dict with keys:
- host, user, password, oauth2_token: connection credentials
- oauth2: dict with provider, client_id, email, client_secret (optional)
Returns:
Healthy IMAP connection, or None if connection failed.
May return a different connection object if reconnection was needed.
"""
if conf.get("oauth2"):
imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
return imap_common.ensure_connection_from_conf(conn, conf)
def ensure_folder_session(conn, conf, folder_name, readonly=True):
"""
Ensure connection is healthy and folder is selected.
Proactively refreshes OAuth2 token if needed. If the connection was
refreshed (new connection object), reselects the folder.
Args:
conn: Existing IMAP connection or None
conf: Connection config dict (see ensure_connection)
folder_name: IMAP folder to select
readonly: Whether to select folder as readonly (default True)
Returns:
Tuple of (connection, success: bool)
- On success: (connection, True) - folder is selected
- On failure: (connection or None, False)
"""
old_conn = conn
new_conn = ensure_connection(conn, conf)
if not new_conn:
return None, False
# If connection changed (token refreshed), need to reselect folder
if new_conn is not old_conn:
try:
new_conn.select(f'"{folder_name}"', readonly=readonly)
except Exception:
return new_conn, False
return new_conn, True

337
src/migrate_imap_emails.py
View File

@ -117,6 +117,7 @@ import threading
import imap_common
import imap_oauth2
import imap_session
# Configuration defaults
DELETE_FROM_SOURCE_DEFAULT = False
@ -288,19 +289,167 @@ def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msg
return deleted_count
def ensure_connection(conn, conf):
"""Refresh OAuth2 token if needed and ensure connection is healthy."""
if conf.get("oauth2"):
imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
return imap_common.ensure_connection_from_conf(conn, conf)
def get_thread_connections(src_conf, dest_conf):
thread_local.src = ensure_connection(getattr(thread_local, "src", None), src_conf)
thread_local.dest = ensure_connection(getattr(thread_local, "dest", None), dest_conf)
thread_local.src = imap_session.ensure_connection(getattr(thread_local, "src", None), src_conf)
thread_local.dest = imap_session.ensure_connection(getattr(thread_local, "dest", None), dest_conf)
return thread_local.src, thread_local.dest
def process_single_uid(
src,
dest,
uid,
folder_name,
delete_from_source,
trash_folder,
preserve_flags,
gmail_mode,
label_index,
check_duplicate,
):
"""
Migrate a single email by UID.
Returns:
Tuple of (success, src, dest, deleted):
- success=True: UID processed (copied, skipped, or non-auth error)
- success=False: Auth error, caller should retry after reconnect
- deleted: 1 if message was marked for deletion, 0 otherwise
"""
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
try:
resp, data = src.uid("fetch", uid, "(FLAGS INTERNALDATE BODY.PEEK[])")
if resp != "OK":
safe_print(f"[{folder_name}] ERROR Fetch | UID {uid_str}")
return (True, src, dest, 0)
msg_content = None
flags = None
date_str = None
for item in data:
if isinstance(item, tuple):
msg_content = item[1]
meta = item[0].decode("utf-8", errors="ignore")
flags_match = re.search(r"FLAGS\s+\((.*?)\)", meta)
if flags_match:
flags = filter_preservable_flags(flags_match.group(1))
date_match = re.search(r"INTERNALDATE\s+\"(.*?)\"", meta)
if date_match:
date_str = f'"{date_match.group(1)}"'
if not msg_content:
return (True, src, dest, 0)
size = len(msg_content) if isinstance(msg_content, (bytes, bytearray)) else 0
size_str = f"{size / 1024:.1f}KB" if size else "0KB"
msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(msg_content)
if not msg_id:
msg_id = imap_common.parse_message_id_from_bytes(msg_content)
# Determine target folder and labels for Gmail mode
apply_labels = gmail_mode
labels = []
if apply_labels and msg_id and label_index is not None:
labels = sorted(label_index.get(msg_id, set()))
if apply_labels:
skip_folders = {imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH}
target_folder = None
remaining_labels = []
for label in labels:
label_folder = label_to_folder(label)
if label_folder in skip_folders:
continue
if target_folder is None:
target_folder = label_folder
else:
remaining_labels.append(label)
if target_folder is None:
target_folder = imap_common.FOLDER_RESTORED_UNLABELED
remaining_labels = []
else:
target_folder = folder_name
imap_common.ensure_folder_exists(dest, target_folder)
dest.select(f'"{target_folder}"')
is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id))
if is_duplicate:
safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}")
if preserve_flags and flags and msg_id:
sync_flags_on_existing(dest, target_folder, msg_id, flags, size)
else:
valid_flags = f"({flags})" if (preserve_flags and flags) else None
imap_common.append_email(
dest,
target_folder,
msg_content,
date_str,
valid_flags,
ensure_folder=False,
)
safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}")
if preserve_flags and flags:
for flag in flags.split():
safe_print(f" -> Applied flag: {flag}")
# Apply remaining Gmail labels
if apply_labels and remaining_labels and msg_id:
for label in remaining_labels:
label_folder = label_to_folder(label)
if label_folder == target_folder:
continue
if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH):
continue
try:
imap_common.ensure_folder_exists(dest, label_folder)
dest.select(f'"{label_folder}"')
if not imap_common.message_exists_in_folder(dest, msg_id):
valid_flags = f"({flags})" if (preserve_flags and flags) else None
imap_common.append_email(
dest,
label_folder,
msg_content,
date_str,
valid_flags,
ensure_folder=False,
)
safe_print(f" -> Applied label: {label}")
if preserve_flags and flags:
for flag in flags.split():
safe_print(f" -> Applied flag: {flag}")
elif preserve_flags and flags:
sync_flags_on_existing(dest, label_folder, msg_id, flags, size)
except Exception as e:
safe_print(f" -> Error applying label {label}: {e}")
deleted = 0
if delete_from_source:
if trash_folder and folder_name != trash_folder:
try:
src.uid("copy", uid, f'"{trash_folder}"')
except Exception:
pass
src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL)
deleted = 1
return (True, src, dest, deleted)
except Exception as e:
if imap_oauth2.is_auth_error(e):
safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
return (False, src, dest, 0)
else:
safe_print(f"[{folder_name}] ERROR Exec | UID {uid_str}: {e}")
return (True, src, dest, 0)
def process_batch(
uids,
folder_name,
@ -318,14 +467,12 @@ def process_batch(
safe_print("Error: Could not establish connections in worker thread.")
return
# Select source folder
try:
src.select(f'"{folder_name}"', readonly=False)
except Exception as e:
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
return
# In non-Gmail-mode, we keep a selected destination folder for efficiency
if not gmail_mode:
try:
imap_common.ensure_folder_exists(dest, folder_name)
@ -335,134 +482,54 @@ def process_batch(
return
deleted_count = 0
for uid in uids:
try:
# Fetch full message (needed to copy and/or apply labels)
resp, data = src.uid("fetch", uid, "(FLAGS INTERNALDATE BODY.PEEK[])")
if resp != "OK":
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
safe_print(f"[{folder_name}] ERROR Fetch | UID {uid_str}")
continue
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
max_retries = 2
msg_content = None
flags = None
date_str = None
for attempt in range(max_retries):
src, src_ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=False)
thread_local.src = src
if not src_ok:
safe_print(f"[{folder_name}] ERROR: Source connection/folder lost for UID {uid_str}")
return
for item in data:
if isinstance(item, tuple):
msg_content = item[1]
meta = item[0].decode("utf-8", errors="ignore")
flags_match = re.search(r"FLAGS\s+\((.*?)\)", meta)
if flags_match:
flags = filter_preservable_flags(flags_match.group(1))
date_match = re.search(r"INTERNALDATE\s+\"(.*?)\"", meta)
if date_match:
date_str = f'"{date_match.group(1)}"'
if not msg_content:
continue
# Compute size and parse headers from the already-fetched message bytes.
size = len(msg_content) if isinstance(msg_content, (bytes, bytearray)) else 0
size_str = f"{size / 1024:.1f}KB" if size else "0KB"
msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(msg_content)
# Prefer Message-ID from body if missing from header parse
if not msg_id:
msg_id = imap_common.parse_message_id_from_bytes(msg_content)
# Determine target folder and labels for Gmail mode
apply_labels = gmail_mode
labels = []
if apply_labels and msg_id and label_index is not None:
labels = sorted(label_index.get(msg_id, set()))
if apply_labels:
skip_folders = {imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH}
target_folder = None
remaining_labels = []
for label in labels:
label_folder = label_to_folder(label)
if label_folder in skip_folders:
continue
if target_folder is None:
target_folder = label_folder
else:
remaining_labels.append(label)
if target_folder is None:
target_folder = imap_common.FOLDER_RESTORED_UNLABELED
remaining_labels = []
if not gmail_mode:
dest, dest_ok = imap_session.ensure_folder_session(dest, dest_conf, folder_name, readonly=False)
thread_local.dest = dest
if not dest_ok:
safe_print(f"[{folder_name}] ERROR: Dest connection/folder lost for UID {uid_str}")
return
else:
target_folder = folder_name
remaining_labels = []
dest = imap_session.ensure_connection(dest, dest_conf)
thread_local.dest = dest
if not dest:
safe_print(f"[{folder_name}] ERROR: Dest connection lost for UID {uid_str}")
return
imap_common.ensure_folder_exists(dest, target_folder)
dest.select(f'"{target_folder}"')
success, src, dest, deleted = process_single_uid(
src,
dest,
uid,
folder_name,
delete_from_source,
trash_folder,
preserve_flags,
gmail_mode,
label_index,
check_duplicate,
)
thread_local.src = src
thread_local.dest = dest
deleted_count += deleted
is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id))
if is_duplicate:
safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}")
if preserve_flags and flags and msg_id:
sync_flags_on_existing(dest, target_folder, msg_id, flags, size)
else:
valid_flags = f"({flags})" if (preserve_flags and flags) else None
imap_common.append_email(
dest,
target_folder,
msg_content,
date_str,
valid_flags,
ensure_folder=False,
)
safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}")
if preserve_flags and flags:
for flag in flags.split():
safe_print(f" -> Applied flag: {flag}")
# Apply remaining Gmail labels (always, whether copied or skipped)
if apply_labels and remaining_labels and msg_id:
for label in remaining_labels:
label_folder = label_to_folder(label)
if label_folder == target_folder:
continue
if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH):
continue
try:
imap_common.ensure_folder_exists(dest, label_folder)
dest.select(f'"{label_folder}"')
if not imap_common.message_exists_in_folder(dest, msg_id):
valid_flags = f"({flags})" if (preserve_flags and flags) else None
imap_common.append_email(
dest,
label_folder,
msg_content,
date_str,
valid_flags,
ensure_folder=False,
)
safe_print(f" -> Applied label: {label}")
if preserve_flags and flags:
for flag in flags.split():
safe_print(f" -> Applied flag: {flag}")
elif preserve_flags and flags:
sync_flags_on_existing(dest, label_folder, msg_id, flags, size)
except Exception as e:
safe_print(f" -> Error applying label {label}: {e}")
if delete_from_source:
if trash_folder and folder_name != trash_folder:
try:
src.uid("copy", uid, f'"{trash_folder}"')
except Exception:
pass
src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL)
deleted_count += 1
except Exception as e:
safe_print(f"[{folder_name}] ERROR Exec | UID {uid}: {e}")
if success:
break
if attempt < max_retries - 1:
src = None
dest = None
thread_local.src = None
thread_local.dest = None
if delete_from_source and deleted_count > 0:
try:
@ -961,11 +1028,11 @@ def main():
safe_print(f"Skipping migration of Trash folder '{name}' (preventing circular migration).")
continue
src_main = ensure_connection(src_main, src_conf)
src_main = imap_session.ensure_connection(src_main, src_conf)
if not src_main:
safe_print("Fatal: Could not reconnect to source IMAP server. Aborting.")
sys.exit(1)
dest_main = ensure_connection(dest_main, dest_conf)
dest_main = imap_session.ensure_connection(dest_main, dest_conf)
if not dest_main:
safe_print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
sys.exit(1)

24
src/restore_imap_emails.py
View File

@ -61,6 +61,7 @@ from typing import Optional
import imap_common
import imap_oauth2
import imap_session
import restore_cache
@ -88,16 +89,9 @@ def safe_print(message):
print(f"[{short_name}] {message}")
def ensure_connection(conn, conf):
"""Refresh OAuth2 token if needed and ensure connection is healthy."""
if conf.get("oauth2"):
imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
return imap_common.ensure_connection_from_conf(conn, conf)
def get_thread_connection(dest_conf):
"""Get or create a thread-local IMAP connection."""
thread_local.dest = ensure_connection(getattr(thread_local, "dest", None), dest_conf)
thread_local.dest = imap_session.ensure_connection(getattr(thread_local, "dest", None), dest_conf)
return thread_local.dest
@ -395,10 +389,16 @@ def process_restore_batch(
gmail_mode = folder_name == "__GMAIL_MODE__"
for file_path, filename in eml_files:
# Proactively refresh token if needed
dest = get_thread_connection(dest_conf)
if not dest:
safe_print(f"ERROR: Connection lost for {filename}")
return
try:
message_id, date_str, raw_content, subject = parse_eml_file(file_path)
if raw_content is None:
continue
continue # No content, skip to next file
size = len(raw_content)
size_str = f"{size / 1024:.1f}KB"
@ -472,7 +472,7 @@ def process_restore_batch(
and message_id in existing_dest_msg_ids
):
safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
continue
continue # Skip to next file
# Upload to target folder.
# Keep server-side duplicate checks enabled to avoid creating duplicates for emails
@ -815,7 +815,7 @@ def restore_folder(
# Delete orphan emails from destination if enabled
if dest_delete and local_msg_ids is not None:
safe_print("Syncing destination: removing emails not in local backup...")
dest = ensure_connection(None, dest_conf)
dest = imap_session.ensure_connection(None, dest_conf)
if dest:
delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids)
dest.logout()
@ -1192,7 +1192,7 @@ def main():
continue
# Proactively refresh OAuth2 token and ensure connection is healthy between folders
dest = ensure_connection(dest, dest_conf)
dest = imap_session.ensure_connection(dest, dest_conf)
if not dest:
print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
sys.exit(1)

16
test/test_backup_imap_emails.py
View File

@ -551,16 +551,16 @@ class TestGmailLabelsPreservation:
"<msg2@test.com>": {"flags": []},
}
def mock_get_message_ids(conn, folder, progress_cb=None):
return folder_data.get(folder, set())
def mock_get_message_info(conn, folder, progress_cb=None):
def mock_get_message_info_with_conf(conn, folder, src_conf, progress_cb=None):
if folder == "[Gmail]/All Mail":
return all_mail_info
return {}
return (all_mail_info, conn)
return ({}, conn)
monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder", mock_get_message_ids)
monkeypatch.setattr(backup_imap_emails, "get_message_info_in_folder", mock_get_message_info)
def mock_get_message_ids_with_conf(conn, folder, src_conf, progress_cb=None):
return (folder_data.get(folder, set()), conn)
monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder_with_conf", mock_get_message_ids_with_conf)
monkeypatch.setattr(backup_imap_emails, "get_message_info_in_folder_with_conf", mock_get_message_info_with_conf)
result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path))

29
test/test_imap_common.py
View File

@ -16,6 +16,8 @@ import os
import sys
from unittest.mock import Mock
import pytest
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
import imap_common
@ -653,14 +655,22 @@ class TestGetMessageIdsInFolder:
result = imap_common.get_message_ids_in_folder(mock_conn)
assert result == {}
def test_search_exception(self):
"""Test returns empty dict when search raises exception."""
def test_search_non_auth_exception_returns_empty(self):
"""Test returns empty dict when search raises non-auth exception."""
mock_conn = Mock()
mock_conn.uid.side_effect = Exception("Connection error")
result = imap_common.get_message_ids_in_folder(mock_conn)
assert result == {}
def test_search_auth_error_reraises(self):
"""Test re-raises auth errors so callers can handle reconnection."""
mock_conn = Mock()
mock_conn.uid.side_effect = Exception("User not authenticated")
with pytest.raises(Exception, match="not authenticated"):
imap_common.get_message_ids_in_folder(mock_conn)
def test_single_message(self):
"""Test fetching single message ID."""
mock_conn = Mock()
@ -707,8 +717,8 @@ class TestGetMessageIdsInFolder:
result = imap_common.get_message_ids_in_folder(mock_conn)
assert result == {}
def test_fetch_exception_for_batch(self):
"""Test continues when fetch raises exception for a batch."""
def test_fetch_non_auth_exception_continues(self):
"""Test continues when fetch raises non-auth exception for a batch."""
mock_conn = Mock()
mock_conn.uid.side_effect = [
("OK", [b"1"]),
@ -718,6 +728,17 @@ class TestGetMessageIdsInFolder:
result = imap_common.get_message_ids_in_folder(mock_conn)
assert result == {}
def test_fetch_auth_error_reraises(self):
"""Test re-raises auth errors during fetch so callers can handle reconnection."""
mock_conn = Mock()
mock_conn.uid.side_effect = [
("OK", [b"1"]),
Exception("AccessTokenExpired"),
]
with pytest.raises(Exception, match="AccessTokenExpired"):
imap_common.get_message_ids_in_folder(mock_conn)
def test_skips_empty_message_id(self):
"""Test that empty message IDs are not added to dict."""
mock_conn = Mock()

100
test/test_imap_oauth2.py
View File

@ -28,9 +28,11 @@ def clear_oauth2_caches():
"""Clear module-level OAuth2 caches between tests."""
imap_oauth2._msal_app_cache.clear()
imap_oauth2._google_creds_cache.clear()
imap_oauth2._tenant_cache.clear()
yield
imap_oauth2._msal_app_cache.clear()
imap_oauth2._google_creds_cache.clear()
imap_oauth2._tenant_cache.clear()
class TestDetectOauth2Provider:
@ -487,3 +489,101 @@ class TestRefreshOauth2Token:
# The other two should see conf["oauth2_token"] changed and skip.
assert call_count["value"] == 1
assert conf["oauth2_token"] == "fresh_token"
class TestIsTokenExpiredError:
"""Tests for is_token_expired_error function."""
def test_detects_access_token_expired(self):
"""Test detects 'accesstokenexpired' error."""
error = Exception("AUTHENTICATE failed: AccessTokenExpired")
assert imap_oauth2.is_token_expired_error(error) is True
def test_detects_session_invalidated(self):
"""Test detects 'session invalidated' error."""
error = Exception("Session invalidated by server")
assert imap_oauth2.is_token_expired_error(error) is True
def test_case_insensitive_access_token_expired(self):
"""Test case-insensitive matching for accesstokenexpired."""
error = Exception("ACCESSTOKENEXPIRED")
assert imap_oauth2.is_token_expired_error(error) is True
def test_case_insensitive_session_invalidated(self):
"""Test case-insensitive matching for session invalidated."""
error = Exception("SESSION INVALIDATED")
assert imap_oauth2.is_token_expired_error(error) is True
def test_false_for_connection_error(self):
"""Test returns False for connection errors."""
error = Exception("Connection refused")
assert imap_oauth2.is_token_expired_error(error) is False
def test_false_for_timeout_error(self):
"""Test returns False for timeout errors."""
error = Exception("The read operation timed out")
assert imap_oauth2.is_token_expired_error(error) is False
def test_false_for_generic_error(self):
"""Test returns False for generic errors."""
error = Exception("Something went wrong")
assert imap_oauth2.is_token_expired_error(error) is False
def test_false_for_not_authenticated(self):
"""Test returns False for 'not authenticated' (handled by is_auth_error)."""
error = Exception("User not authenticated")
assert imap_oauth2.is_token_expired_error(error) is False
class TestIsAuthError:
"""Tests for is_auth_error function."""
def test_detects_access_token_expired(self):
"""Test detects 'accesstokenexpired' error (via is_token_expired_error)."""
error = Exception("AUTHENTICATE failed: AccessTokenExpired")
assert imap_oauth2.is_auth_error(error) is True
def test_detects_session_invalidated(self):
"""Test detects 'session invalidated' error (via is_token_expired_error)."""
error = Exception("Session invalidated by server")
assert imap_oauth2.is_auth_error(error) is True
def test_detects_not_authenticated(self):
"""Test detects 'not authenticated' error."""
error = Exception("User not authenticated")
assert imap_oauth2.is_auth_error(error) is True
def test_detects_authentication_failed(self):
"""Test detects 'authentication failed' error."""
error = Exception("Authentication failed for user@example.com")
assert imap_oauth2.is_auth_error(error) is True
def test_case_insensitive_not_authenticated(self):
"""Test case-insensitive matching for not authenticated."""
error = Exception("NOT AUTHENTICATED")
assert imap_oauth2.is_auth_error(error) is True
def test_case_insensitive_authentication_failed(self):
"""Test case-insensitive matching for authentication failed."""
error = Exception("AUTHENTICATION FAILED")
assert imap_oauth2.is_auth_error(error) is True
def test_false_for_connection_error(self):
"""Test returns False for connection errors."""
error = Exception("Connection refused")
assert imap_oauth2.is_auth_error(error) is False
def test_false_for_timeout_error(self):
"""Test returns False for timeout errors."""
error = Exception("The read operation timed out")
assert imap_oauth2.is_auth_error(error) is False
def test_false_for_generic_error(self):
"""Test returns False for generic errors."""
error = Exception("Something went wrong")
assert imap_oauth2.is_auth_error(error) is False
def test_false_for_folder_not_found(self):
"""Test returns False for folder errors."""
error = Exception("Folder not found: INBOX")
assert imap_oauth2.is_auth_error(error) is False

268
test/test_imap_session.py Normal file
View File

@ -0,0 +1,268 @@
"""
Tests for imap_session.py
Tests cover:
- ensure_connection() behavior with and without OAuth2
- ensure_folder_session() connection change detection
- Folder selection and re-selection logic
- Error handling for connection and folder failures
"""
import os
import sys
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
import imap_session
class TestEnsureConnection:
"""Tests for ensure_connection function."""
def test_with_oauth2_calls_refresh(self):
"""Test that OAuth2 token refresh is called when oauth2 config is present."""
conf = {
"host": "imap.example.com",
"user": "user@example.com",
"password": "pass",
"oauth2_token": "old_token",
"oauth2": {
"provider": "microsoft",
"client_id": "client-id",
"email": "user@example.com",
},
}
mock_conn = MagicMock()
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
with patch.object(
imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn
) as mock_ensure:
result = imap_session.ensure_connection(mock_conn, conf)
mock_refresh.assert_called_once_with(conf, "old_token")
mock_ensure.assert_called_once_with(mock_conn, conf)
assert result is mock_conn
def test_without_oauth2_skips_refresh(self):
"""Test that token refresh is skipped when oauth2 config is not present."""
conf = {
"host": "imap.example.com",
"user": "user@example.com",
"password": "pass",
"oauth2": None,
}
mock_conn = MagicMock()
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
imap_session.ensure_connection(mock_conn, conf)
mock_refresh.assert_not_called()
def test_without_oauth2_key_skips_refresh(self):
"""Test that token refresh is skipped when oauth2 key is missing entirely."""
conf = {
"host": "imap.example.com",
"user": "user@example.com",
"password": "pass",
}
mock_conn = MagicMock()
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
imap_session.ensure_connection(mock_conn, conf)
mock_refresh.assert_not_called()
def test_returns_connection_from_ensure_connection_from_conf(self):
"""Test returns the connection from ensure_connection_from_conf."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn):
result = imap_session.ensure_connection(old_conn, conf)
assert result is new_conn
def test_returns_none_when_connection_fails(self):
"""Test returns None when ensure_connection_from_conf fails."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=None):
result = imap_session.ensure_connection(MagicMock(), conf)
assert result is None
def test_with_none_connection(self):
"""Test works correctly when passed None as connection."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
new_conn = MagicMock()
with patch.object(
imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn
) as mock_ensure:
result = imap_session.ensure_connection(None, conf)
mock_ensure.assert_called_once_with(None, conf)
assert result is new_conn
class TestEnsureFolderSession:
"""Tests for ensure_folder_session function."""
def test_same_connection_no_folder_select(self):
"""Test that folder is not re-selected when connection stays the same."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
mock_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=mock_conn):
result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
assert result_conn is mock_conn
assert success is True
mock_conn.select.assert_not_called()
def test_new_connection_selects_folder(self):
"""Test that folder is selected when connection changes."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
def test_new_connection_selects_folder_readwrite(self):
"""Test that folder is selected with readonly=False when specified."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(
old_conn, conf, "[Gmail]/All Mail", readonly=False
)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"[Gmail]/All Mail"', readonly=False)
def test_none_connection_returns_failure(self):
"""Test returns (None, False) when ensure_connection returns None."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
mock_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=None):
result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
assert result_conn is None
assert success is False
def test_folder_select_failure_returns_false(self):
"""Test returns (conn, False) when folder selection fails."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
new_conn.select.side_effect = Exception("Folder not found")
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "NonExistent", readonly=True)
assert result_conn is new_conn
assert success is False
def test_none_old_connection_selects_folder(self):
"""Test that folder is selected when old connection was None."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(None, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
def test_folder_name_with_spaces(self):
"""Test folder selection with folder names containing spaces."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "My Folder", readonly=True)
assert success is True
new_conn.select.assert_called_once_with('"My Folder"', readonly=True)
def test_folder_select_imap_error(self):
"""Test handles IMAP-specific errors during folder selection."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
new_conn.select.side_effect = Exception("IMAP error: NO SELECT failed")
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is False
class TestEnsureFolderSessionWithOAuth2:
"""Integration-style tests for ensure_folder_session with OAuth2."""
def test_oauth2_refresh_triggers_folder_reselect(self):
"""Test that OAuth2 token refresh (new connection) triggers folder reselection."""
conf = {
"host": "outlook.office365.com",
"user": "user@example.com",
"password": "oauth2_token_here",
"oauth2_token": "old_token",
"oauth2": {
"provider": "microsoft",
"client_id": "client-id",
"email": "user@example.com",
},
}
old_conn = MagicMock()
new_conn = MagicMock()
# Simulate token refresh causing a new connection
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"):
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
def test_healthy_connection_no_reselect(self):
"""Test that healthy connection (same object) doesn't reselect folder."""
conf = {
"host": "outlook.office365.com",
"user": "user@example.com",
"password": "oauth2_token_here",
"oauth2_token": "valid_token",
"oauth2": {
"provider": "microsoft",
"client_id": "client-id",
"email": "user@example.com",
},
}
mock_conn = MagicMock()
# Simulate healthy connection (same object returned)
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"):
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
assert result_conn is mock_conn
assert success is True
mock_conn.select.assert_not_called()