Add imap_session module for centralized OAuth2 refresh (#21)
* Add .claude to gitignore * Add imap_session module for centralized OAuth2 refresh Create new imap_session.py module that combines imap_common and imap_oauth2 for unified session management: - ensure_connection(): Refreshes OAuth2 token and ensures connection health - ensure_folder_session(): Ensures connection AND folder selection Updated backup, migrate, and restore scripts to use the new module, simplifying the proactive token refresh pattern in process_batch functions. * Move is_token_expired_error to imap_oauth2 module Relocates the OAuth2 token expiration detection function to the oauth2-specific module where it logically belongs. * Add retry logic for auth errors during email processing Extract single-UID processing into dedicated functions with retry capability for recoverable authentication failures. Broadens auth error detection beyond token expiration to include session loss and authentication failures. * Cache Microsoft tenant discovery results Avoids repeated OpenID Connect discovery requests for the same email domain during OAuth2 authentication flows. Normalizes domain to lowercase for consistent cache keys. * Add tests for is_token_expired_error and is_auth_error Covers token expiration, session invalidation, and authentication failure detection with case-insensitive matching verification. Ensures non-auth errors like connection and timeout issues are correctly excluded. * Add tests for imap_session module Covers ensure_connection and ensure_folder_session functions including OAuth2 token refresh, folder selection, and error handling scenarios. Verifies connection change detection triggers folder re-selection. * Add tests for auth error re-raising in get_message_ids_in_folder Verifies that authentication errors are propagated to callers for reconnection handling while non-auth errors return empty results. Clarifies test names to distinguish error handling behaviors.
This commit is contained in:
Nathan Moinvaziri
GitHub
parent
8c49bc7c48
commit
e932190618
1
.gitignore
vendored
1
.gitignore
vendored
@ -7,3 +7,4 @@ env/
|
||||
venv/
|
||||
.env
|
||||
.coverage
|
||||
.claude
|
||||
|
||||
@ -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)
|
||||
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]
|
||||
success, src = process_single_uid(src, uid, folder_name, local_folder_path)
|
||||
thread_local.src = src
|
||||
|
||||
if success:
|
||||
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)
|
||||
# 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 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:
|
||||
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)
|
||||
|
||||
try:
|
||||
src.logout()
|
||||
except Exception:
|
||||
pass # Connection may already be closed
|
||||
print("\nBackup completed successfully.")
|
||||
|
||||
if args.preserve_labels or args.gmail_mode:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
62
src/imap_session.py
Normal 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
|
||||
@ -117,6 +117,7 @@ import threading
|
||||
|
||||
import imap_common
|
||||
import imap_oauth2
|
||||
import imap_session
|
||||
|
||||
# Configuration defaults
|
||||
DELETE_FROM_SOURCE_DEFAULT = False
|
||||
@ -288,61 +289,40 @@ 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_batch(
|
||||
uids,
|
||||
def process_single_uid(
|
||||
src,
|
||||
dest,
|
||||
uid,
|
||||
folder_name,
|
||||
src_conf,
|
||||
dest_conf,
|
||||
delete_from_source,
|
||||
trash_folder=None,
|
||||
preserve_flags=False,
|
||||
gmail_mode=False,
|
||||
label_index=None,
|
||||
check_duplicate=True,
|
||||
trash_folder,
|
||||
preserve_flags,
|
||||
gmail_mode,
|
||||
label_index,
|
||||
check_duplicate,
|
||||
):
|
||||
src, dest = get_thread_connections(src_conf, dest_conf)
|
||||
if not src or not dest:
|
||||
safe_print("Error: Could not establish connections in worker thread.")
|
||||
return
|
||||
"""
|
||||
Migrate a single email by UID.
|
||||
|
||||
# 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
|
||||
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)
|
||||
|
||||
# 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)
|
||||
dest.select(f'"{folder_name}"')
|
||||
except Exception as e:
|
||||
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
|
||||
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
|
||||
return (True, src, dest, 0)
|
||||
|
||||
msg_content = None
|
||||
flags = None
|
||||
@ -360,14 +340,12 @@ def process_batch(
|
||||
date_str = f'"{date_match.group(1)}"'
|
||||
|
||||
if not msg_content:
|
||||
continue
|
||||
return (True, src, dest, 0)
|
||||
|
||||
# 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)
|
||||
|
||||
@ -396,7 +374,6 @@ def process_batch(
|
||||
remaining_labels = []
|
||||
else:
|
||||
target_folder = folder_name
|
||||
remaining_labels = []
|
||||
|
||||
imap_common.ensure_folder_exists(dest, target_folder)
|
||||
dest.select(f'"{target_folder}"')
|
||||
@ -422,7 +399,7 @@ def process_batch(
|
||||
for flag in flags.split():
|
||||
safe_print(f" -> Applied flag: {flag}")
|
||||
|
||||
# Apply remaining Gmail labels (always, whether copied or skipped)
|
||||
# Apply remaining Gmail labels
|
||||
if apply_labels and remaining_labels and msg_id:
|
||||
for label in remaining_labels:
|
||||
label_folder = label_to_folder(label)
|
||||
@ -452,6 +429,7 @@ def process_batch(
|
||||
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:
|
||||
@ -459,10 +437,99 @@ def process_batch(
|
||||
except Exception:
|
||||
pass
|
||||
src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL)
|
||||
deleted_count += 1
|
||||
deleted = 1
|
||||
|
||||
return (True, src, dest, deleted)
|
||||
|
||||
except Exception as e:
|
||||
safe_print(f"[{folder_name}] ERROR Exec | UID {uid}: {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,
|
||||
src_conf,
|
||||
dest_conf,
|
||||
delete_from_source,
|
||||
trash_folder=None,
|
||||
preserve_flags=False,
|
||||
gmail_mode=False,
|
||||
label_index=None,
|
||||
check_duplicate=True,
|
||||
):
|
||||
src, dest = get_thread_connections(src_conf, dest_conf)
|
||||
if not src or not dest:
|
||||
safe_print("Error: Could not establish connections in worker thread.")
|
||||
return
|
||||
|
||||
try:
|
||||
src.select(f'"{folder_name}"', readonly=False)
|
||||
except Exception as e:
|
||||
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
|
||||
return
|
||||
|
||||
if not gmail_mode:
|
||||
try:
|
||||
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}")
|
||||
return
|
||||
|
||||
deleted_count = 0
|
||||
|
||||
for uid in uids:
|
||||
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
|
||||
max_retries = 2
|
||||
|
||||
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
|
||||
|
||||
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:
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
268
test/test_imap_session.py
Normal 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()
|
||||
Loading…
Reference in New Issue
Block a user
