diff --git a/src/imap_common.py b/src/imap_common.py index 5447f92..7d03822 100644 --- a/src/imap_common.py +++ b/src/imap_common.py @@ -248,45 +248,6 @@ def parse_message_id_and_subject_from_bytes(raw_message): return None, "(No Subject)" -def get_msg_details(imap_conn, uid): - """ - Fetches simplified message details (Message-ID, Size, Subject) for a given UID. - Returns (msg_id, size, subject) tuple. - """ - try: - resp, data = imap_conn.uid("fetch", uid, "(RFC822.SIZE BODY.PEEK[HEADER.FIELDS (MESSAGE-ID SUBJECT)])") - except Exception: - return None, None, None - - if resp != "OK": - return None, None, None - - msg_id = None - subject = "(No Subject)" - size = 0 - - for item in data: - if isinstance(item, tuple): - content = item[0].decode("utf-8", errors="ignore") - - # Parse Size - size_match = re.search(r"RFC822\.SIZE\s+(\d+)", content) - if size_match: - size = int(size_match.group(1)) - - # Parse Headers - msg_bytes = item[1] - # Use compat32 to preserve raw headers with continuation lines - parser = BytesParser(policy=policy.compat32) - email_obj = parser.parsebytes(msg_bytes, headersonly=True) - msg_id = decode_message_id(email_obj.get("Message-ID")) - raw_subject = email_obj.get("Subject") - if raw_subject: - subject = decode_mime_header(raw_subject) - - return msg_id, size, subject - - def message_exists_in_folder(dest_conn, msg_id): """ Checks if a message with the given Message-ID exists in the CURRENTLY SELECTED folder of dest_conn. @@ -307,6 +268,63 @@ def message_exists_in_folder(dest_conn, msg_id): return False +def get_message_ids_in_folder(imap_conn): + """ + Fetches all Message-IDs from the currently selected folder. + Returns a dict mapping UID (bytes) -> Message-ID (str). + UIDs without a Message-ID are not included in the result. + Use set(result.values()) to get just the Message-ID set. + """ + try: + resp, data = imap_conn.uid("search", None, "ALL") + if resp != "OK" or not data[0].strip(): + return {} + except Exception: + return {} + + uids = data[0].split() + return get_uid_to_message_id_map(imap_conn, uids) + + +def get_uid_to_message_id_map(imap_conn, uids): + """ + Fetches Message-IDs for a list of UIDs from the currently selected folder. + Returns a dict mapping UID (bytes) -> Message-ID (str). + UIDs without a Message-ID are not included in the result. + """ + uid_to_msgid = {} + if not uids: + return uid_to_msgid + + FETCH_BATCH = 500 + for i in range(0, len(uids), FETCH_BATCH): + batch = uids[i : i + FETCH_BATCH] + uid_range = b",".join(batch) + try: + resp, fetch_data = imap_conn.uid("fetch", uid_range, "(UID BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])") + if resp != "OK": + continue + for item in fetch_data: + if isinstance(item, tuple) and len(item) >= 2: + # Extract UID from response like b'1 (UID 12345 BODY[HEADER.FIELDS ...]' + meta = item[0] + if isinstance(meta, bytes): + meta = meta.decode("utf-8", errors="ignore") + uid_match = re.search(r"UID\s+(\d+)", meta) + if not uid_match: + continue + uid = uid_match.group(1).encode() + + # Extract Message-ID + mid = extract_message_id(item[1]) + if mid: + uid_to_msgid[uid] = mid + except Exception: + continue + + return uid_to_msgid + + def sanitize_filename(filename): """ Sanitizes a string to be safe for use as a filename. diff --git a/src/migrate_imap_emails.py b/src/migrate_imap_emails.py index d927c5b..4770f22 100644 --- a/src/migrate_imap_emails.py +++ b/src/migrate_imap_emails.py @@ -233,7 +233,12 @@ def build_gmail_label_index(src_conn): total = len(label_folders) for i, folder in enumerate(label_folders, start=1): safe_print(f"[{i}/{total}] Scanning label folder for Message-IDs: {folder}") - msg_ids = get_message_ids_in_folder(src_conn, folder) + try: + src_conn.select(f'"{folder}"', readonly=True) + msg_ids = set(imap_common.get_message_ids_in_folder(src_conn).values()) + except Exception as e: + safe_print(f"Error getting message IDs from {folder}: {e}") + msg_ids = set() label = folder_to_label(folder) for msg_id in msg_ids: label_index.setdefault(msg_id, set()).add(label) @@ -241,93 +246,29 @@ def build_gmail_label_index(src_conn): return label_index -def get_message_ids_in_folder(imap_conn, folder_name): - """ - Get a set of Message-IDs for all emails in a folder. - Used for destination deletion sync. - """ - message_ids = set() - try: - imap_conn.select(f'"{folder_name}"', readonly=True) - resp, data = imap_conn.uid("search", None, "ALL") - if resp != "OK" or not data or not data[0]: - return message_ids - - uids = data[0].split() - if not uids: - return message_ids - - # Fetch Message-IDs in batches - batch_size = 200 - for i in range(0, len(uids), batch_size): - batch = uids[i : i + batch_size] - uid_range = b",".join(batch) - try: - resp, items = imap_conn.uid("fetch", uid_range, "(BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])") - if resp != "OK": - continue - - for item in items: - if isinstance(item, tuple) and len(item) >= 2: - msg_id = imap_common.extract_message_id(item[1]) - if msg_id: - message_ids.add(msg_id) - except Exception: - continue - except Exception as e: - safe_print(f"Error getting message IDs from {folder_name}: {e}") - - return message_ids - - -def delete_orphan_emails(imap_conn, folder_name, source_msg_ids): +def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msgid=None): """ Delete emails from destination folder that don't exist in source. Returns count of deleted emails. + + If dest_uid_to_msgid is provided (dict of UID -> Message-ID), it will be used + instead of fetching from the server, avoiding redundant IMAP calls. """ deleted_count = 0 try: imap_conn.select(f'"{folder_name}"', readonly=False) - resp, data = imap_conn.uid("search", None, "ALL") - if resp != "OK" or not data or not data[0]: - return 0 - uids = data[0].split() - if not uids: - return 0 + # Use provided map or fetch from server + if dest_uid_to_msgid is None: + dest_uid_to_msgid = imap_common.get_message_ids_in_folder(imap_conn) - # Check each UID's Message-ID against source - batch_size = 100 + # Find UIDs to delete (in destination but not in source) uids_to_delete = [] - - for i in range(0, len(uids), batch_size): - batch = uids[i : i + batch_size] - uid_range = b",".join(batch) - try: - resp, items = imap_conn.uid("fetch", uid_range, "(UID BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])") - if resp != "OK": - continue - - for item in items: - if isinstance(item, tuple) and len(item) >= 2: - # Extract UID from response - meta_str = ( - item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0]) - ) - uid_match = re.search(r"UID\s+(\d+)", meta_str) - if not uid_match: - continue - uid = uid_match.group(1) - - # Extract Message-ID - msg_id = imap_common.extract_message_id(item[1]) - - # If not in source, mark for deletion - if msg_id and msg_id not in source_msg_ids: - uids_to_delete.append(uid) - - except Exception: - continue + for uid, msg_id in dest_uid_to_msgid.items(): + if msg_id not in source_msg_ids: + # Convert bytes UID to string for STORE command + uid_str = uid.decode() if isinstance(uid, bytes) else str(uid) + uids_to_delete.append(uid_str) # Delete orphan emails for uid in uids_to_delete: @@ -389,6 +330,7 @@ def process_batch( 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: @@ -485,7 +427,7 @@ def process_batch( ensure_dest_folder(target_folder) dest.select(f'"{target_folder}"') - is_duplicate = bool(msg_id and imap_common.message_exists_in_folder(dest, msg_id)) + 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]}") @@ -591,23 +533,54 @@ def migrate_folder( delete_orphan_emails(dest, folder_name, set()) return - safe_print(f"Found {total} messages. Starting parallel migration...") + # Pre-fetch destination Message-IDs for fast duplicate detection (non-Gmail mode only) + dest_msg_ids = None + if not gmail_mode: + safe_print(f"Pre-fetching destination Message-IDs for {folder_name}...") + dest_uid_to_msgid = imap_common.get_message_ids_in_folder(dest) + dest_msg_ids = set(dest_uid_to_msgid.values()) + safe_print(f"Found {len(dest_msg_ids)} existing messages in destination.") - # If dest_delete is enabled, gather source Message-IDs first - source_msg_ids = None - if dest_delete and not gmail_mode: - safe_print("Building source Message-ID index for sync...") - source_msg_ids = get_message_ids_in_folder(src, folder_name) - safe_print(f"Found {len(source_msg_ids)} unique Message-IDs in source.") + # Pre-fetch source Message-IDs and filter out duplicates before processing + # Skip pre-filtering when preserve_flags is True (need to sync flags on duplicates) + uids_to_process = uids + src_msg_ids = None + skipped_duplicate_uids = [] + pre_filtered = False + if not gmail_mode and dest_msg_ids is not None and not preserve_flags: + pre_filtered = True + safe_print(f"Pre-fetching source Message-IDs for {folder_name}...") + # Use get_uid_to_message_id_map directly since we already have UIDs from folder select + src_uid_to_msgid = imap_common.get_uid_to_message_id_map(src, uids) + src_msg_ids = set(src_uid_to_msgid.values()) + + # Filter to only UIDs that need migration (not already in destination) + uids_to_process = [] + for uid in uids: + msg_id = src_uid_to_msgid.get(uid) + if msg_id not in dest_msg_ids: + uids_to_process.append(uid) + else: + skipped_duplicate_uids.append(uid) + safe_print(f"Skipping {len(skipped_duplicate_uids)} duplicates, {len(uids_to_process)} to migrate.") elif dest_delete and gmail_mode: safe_print("Warning: --dest-delete is not supported in --gmail-mode; ignoring.") - # Create batches - uid_batches = [uids[i : i + BATCH_SIZE] for i in range(0, len(uids), BATCH_SIZE)] + if not uids_to_process: + safe_print(f"No new messages to migrate in {folder_name}.") + if dest_delete and src_msg_ids is not None: + safe_print("Syncing destination: removing emails not in source...") + delete_orphan_emails(dest, folder_name, src_msg_ids, dest_uid_to_msgid) + return + + # Create batches from filtered UIDs + uid_batches = [uids_to_process[i : i + BATCH_SIZE] for i in range(0, len(uids_to_process), BATCH_SIZE)] executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) try: futures = [] + # When pre-filtered, we know UIDs are non-duplicates; otherwise need to check + check_duplicate = not pre_filtered for batch in uid_batches: futures.append( executor.submit( @@ -621,6 +594,7 @@ def migrate_folder( preserve_flags, gmail_mode, label_index, + check_duplicate, ) ) @@ -638,7 +612,17 @@ def migrate_folder( executor.shutdown(wait=True) if delete_from_source: - safe_print(f"Expunging any remaining deleted messages from {folder_name}...") + # Delete skipped duplicates from source (already exist in destination) + if skipped_duplicate_uids: + safe_print(f"Deleting {len(skipped_duplicate_uids)} duplicates from source...") + try: + src.select(f'"{folder_name}"', readonly=False) + for uid in skipped_duplicate_uids: + src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL) + except Exception as e: + safe_print(f"Error deleting duplicates: {e}") + + safe_print(f"Expunging deleted messages from {folder_name}...") try: src.select(f'"{folder_name}"', readonly=False) src.expunge() @@ -646,9 +630,9 @@ def migrate_folder( safe_print(f"Error Expunging: {e}") # Delete orphan emails from destination if enabled - if dest_delete and source_msg_ids is not None and not gmail_mode: + if dest_delete and src_msg_ids is not None and not gmail_mode: safe_print("Syncing destination: removing emails not in source...") - delete_orphan_emails(dest, folder_name, source_msg_ids) + delete_orphan_emails(dest, folder_name, src_msg_ids, dest_uid_to_msgid) def main(): diff --git a/src/restore_imap_emails.py b/src/restore_imap_emails.py index 85e0e59..aad0aaf 100644 --- a/src/restore_imap_emails.py +++ b/src/restore_imap_emails.py @@ -206,6 +206,22 @@ def sync_flags_on_existing(imap_conn, folder_name, message_id, flags, size): safe_print(f" -> Error syncing flags: {e}") +def extract_message_id_from_eml(file_path): + """ + Extract just the Message-ID from an .eml file efficiently. + Returns the Message-ID string or None on error. + """ + try: + with open(file_path, "rb") as f: + # Read just the first 64KB to get headers + header_bytes = f.read(65536) + + msg_id = imap_common.extract_message_id(header_bytes) + return msg_id + except Exception: + return None + + def parse_eml_file(file_path): """ Parse an .eml file and extract metadata. @@ -215,11 +231,13 @@ def parse_eml_file(file_path): with open(file_path, "rb") as f: raw_content = f.read() - parser = BytesParser(policy=policy.default) + # Use compat32 to preserve raw headers with continuation lines + parser = BytesParser(policy=policy.compat32) msg = parser.parsebytes(raw_content, headersonly=True) - message_id = msg.get("Message-ID", "").strip() - subject = msg.get("Subject", "(No Subject)") + message_id = imap_common.decode_message_id(msg.get("Message-ID")) + raw_subject = msg.get("Subject") + subject = imap_common.decode_mime_header(raw_subject) if raw_subject else "(No Subject)" date_header = msg.get("Date") # Parse date for IMAP INTERNALDATE @@ -272,13 +290,14 @@ def email_exists_in_folder(imap_conn, message_id): return False -def upload_email(dest, folder_name, raw_content, date_str, message_id, subject, flags=None): +def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True): """ Upload a single email to the destination folder. - Returns True on success, False on failure. + Returns True on success, False if duplicate or on failure. Args: flags: Optional string of IMAP flags like "\\Seen" for read emails. + check_duplicate: Whether to check for duplicates before uploading. """ try: # Ensure folder exists @@ -291,9 +310,8 @@ def upload_email(dest, folder_name, raw_content, date_str, message_id, subject, # Select folder dest.select(f'"{folder_name}"') - # Check for duplicates - size = len(raw_content) - if message_id and email_exists_in_folder(dest, message_id): + # Check for duplicates if requested + if check_duplicate and message_id and email_exists_in_folder(dest, message_id): return False # Already exists # Upload with original date and flags @@ -397,8 +415,12 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab target_folder = folder_name remaining_labels = labels - # Upload to target folder (or check if exists) - uploaded = upload_email(dest, target_folder, raw_content, date_str, message_id, display_subject, flags) + # Upload to target folder + # In Gmail mode, check duplicates per-email since target folders vary + # In non-Gmail mode, duplicates were pre-filtered so no check needed + uploaded = upload_email( + dest, target_folder, raw_content, date_str, message_id, flags, check_duplicate=gmail_mode + ) if not uploaded: safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}") @@ -562,8 +584,50 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la local_msg_ids = get_local_message_ids(local_folder_path) safe_print(f"Found {len(local_msg_ids)} unique Message-IDs in local backup.") + # Pre-fetch destination Message-IDs and filter duplicates (non-Gmail mode only) + gmail_mode = folder_name == "__GMAIL_MODE__" + files_to_restore = eml_files + + if not gmail_mode: + dest_msg_ids = set() + try: + dest_tmp = imap_common.get_imap_connection_from_conf(dest_conf) + if dest_tmp: + # Ensure folder exists before selecting + if folder_name.upper() != "INBOX": + try: + dest_tmp.create(f'"{folder_name}"') + except Exception: + pass + dest_tmp.select(f'"{folder_name}"') + dest_msg_ids = set(imap_common.get_message_ids_in_folder(dest_tmp).values()) + dest_tmp.logout() + except Exception: + dest_msg_ids = set() + + safe_print(f"{len(dest_msg_ids)} existing messages in destination.") + + # Pre-filter files to skip duplicates + if dest_msg_ids: + safe_print("Pre-filtering duplicates...") + files_to_restore = [] + skipped = 0 + for file_path, filename in eml_files: + msg_id = extract_message_id_from_eml(file_path) + if msg_id and msg_id in dest_msg_ids: + skipped += 1 + else: + files_to_restore.append((file_path, filename)) + safe_print(f"Skipping {skipped} duplicates, {len(files_to_restore)} to restore.") + + if not files_to_restore: + safe_print("No new emails to restore.") + return + + safe_print(f"Starting parallel restore of {len(files_to_restore)} emails...") + # Create batches - batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)] + batches = [files_to_restore[i : i + BATCH_SIZE] for i in range(0, len(files_to_restore), BATCH_SIZE)] with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: futures = [] diff --git a/test/test_backup_imap_emails.py b/test/test_backup_imap_emails.py index b179f10..3d88a10 100644 --- a/test/test_backup_imap_emails.py +++ b/test/test_backup_imap_emails.py @@ -287,8 +287,6 @@ class TestBackupErrorHandling: """Test handling of fetch body failure.""" mock_conn = MagicMock() mock_conn.select.return_value = "OK" - # get_msg_details calls fetch headers, make it work - monkeypatch.setattr("imap_common.get_msg_details", lambda conn, uid: (uid, 100, "Subject")) # Fetch body fails mock_conn.uid.return_value = ("NO", [None]) @@ -307,7 +305,6 @@ class TestBackupErrorHandling: monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn) # Mock fetched data - monkeypatch.setattr("imap_common.get_msg_details", lambda conn, uid: (uid, 100, "Subject")) mock_conn.uid.return_value = ("OK", [(b"1 (RFC822 {10}", b"Content")]) # Mock open to fail diff --git a/test/test_imap_common.py b/test/test_imap_common.py index 034fb17..6d9e6a5 100644 --- a/test/test_imap_common.py +++ b/test/test_imap_common.py @@ -310,30 +310,6 @@ class TestMessageExistsInFolder: assert result is False -class TestGetMsgDetails: - """Tests for get_msg_details function.""" - - def test_fetch_error(self): - """Test returns None tuple on fetch error.""" - mock_conn = Mock() - mock_conn.uid.side_effect = Exception("Fetch error") - - msg_id, size, subject = imap_common.get_msg_details(mock_conn, b"1") - assert msg_id is None - assert size is None - assert subject is None - - def test_not_ok_response(self): - """Test returns None tuple on non-OK response.""" - mock_conn = Mock() - mock_conn.uid.return_value = ("NO", None) - - msg_id, size, subject = imap_common.get_msg_details(mock_conn, b"1") - assert msg_id is None - assert size is None - assert subject is None - - class TestExtractMessageId: """Tests for extract_message_id function.""" @@ -372,3 +348,107 @@ class TestExtractMessageId: """Test empty input.""" assert imap_common.extract_message_id(None) is None assert imap_common.extract_message_id(b"") is None + + +class TestGetMessageIdsInFolder: + """Tests for get_message_ids_in_folder function.""" + + def test_empty_folder(self): + """Test returns empty dict for folder with no messages.""" + mock_conn = Mock() + mock_conn.uid.return_value = ("OK", [b""]) + + result = imap_common.get_message_ids_in_folder(mock_conn) + assert result == {} + + def test_search_fails(self): + """Test returns empty dict when search fails.""" + mock_conn = Mock() + mock_conn.uid.return_value = ("NO", []) + + 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.""" + 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_single_message(self): + """Test fetching single message ID.""" + mock_conn = Mock() + # First call: search returns one UID + # Second call: fetch returns the Message-ID header (with UID in response) + mock_conn.uid.side_effect = [ + ("OK", [b"1"]), + ("OK", [(b"1 (UID 1 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"), b")"]), + ] + + result = imap_common.get_message_ids_in_folder(mock_conn) + assert result == {b"1": ""} + assert set(result.values()) == {""} + + def test_multiple_messages(self): + """Test fetching multiple message IDs.""" + mock_conn = Mock() + mock_conn.uid.side_effect = [ + ("OK", [b"1 2 3"]), + ( + "OK", + [ + (b"1 (UID 1 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"), + b")", + (b"2 (UID 2 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"), + b")", + (b"3 (UID 3 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"), + b")", + ], + ), + ] + + result = imap_common.get_message_ids_in_folder(mock_conn) + assert set(result.values()) == {"", "", ""} + + def test_fetch_fails_for_batch(self): + """Test continues when fetch fails for a batch.""" + mock_conn = Mock() + mock_conn.uid.side_effect = [ + ("OK", [b"1"]), + ("NO", []), # Fetch fails + ] + + 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.""" + mock_conn = Mock() + mock_conn.uid.side_effect = [ + ("OK", [b"1"]), + Exception("Fetch error"), + ] + + result = imap_common.get_message_ids_in_folder(mock_conn) + assert result == {} + + def test_skips_empty_message_id(self): + """Test that empty message IDs are not added to dict.""" + mock_conn = Mock() + mock_conn.uid.side_effect = [ + ("OK", [b"1 2"]), + ( + "OK", + [ + (b"1 (UID 1 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"), + b")", + (b"2 (UID 2 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"), # Empty + b")", + ], + ), + ] + + result = imap_common.get_message_ids_in_folder(mock_conn) + assert set(result.values()) == {""} diff --git a/test/test_migrate_imap_emails.py b/test/test_migrate_imap_emails.py index fccc1fb..39294a6 100644 --- a/test/test_migrate_imap_emails.py +++ b/test/test_migrate_imap_emails.py @@ -733,41 +733,3 @@ class TestDestDeleteFunctionality: # All dest emails should be deleted assert len(dest_server.folders["INBOX"]) == 0 - - def test_get_message_ids_in_folder(self, mock_server_factory, monkeypatch): - """Test get_message_ids_in_folder returns correct Message-IDs.""" - data = { - "INBOX": [ - b"Subject: Email 1\r\nMessage-ID: \r\n\r\nBody", - b"Subject: Email 2\r\nMessage-ID: \r\n\r\nBody", - b"Subject: Email 3\r\nMessage-ID: \r\n\r\nBody", - ] - } - - server, port = None, None - # Use single_mock_server fixture pattern - import time - - from conftest import get_free_port, start_server_thread - - port = get_free_port() - thread, server = start_server_thread(port, data) - time.sleep(0.3) - - try: - import imaplib - - conn = imaplib.IMAP4("localhost", port) - conn.login("user", "pass") - - msg_ids = migrate_imap_emails.get_message_ids_in_folder(conn, "INBOX") - - assert "" in msg_ids - assert "" in msg_ids - assert "" in msg_ids - assert len(msg_ids) == 3 - - conn.logout() - finally: - server.shutdown() - thread.join(timeout=2) diff --git a/test/test_restore_imap_emails.py b/test/test_restore_imap_emails.py index a19e26a..b4d6087 100644 --- a/test/test_restore_imap_emails.py +++ b/test/test_restore_imap_emails.py @@ -136,7 +136,7 @@ Body content. message_id, date_str, raw_content, subject = restore_imap_emails.parse_eml_file(str(eml_file)) - assert message_id == "" + assert message_id is None assert raw_content is not None def test_parse_nonexistent_file(self): @@ -273,8 +273,8 @@ class TestUploadEmail: mock_conn.select.return_value = ("OK", [b"1"]) mock_conn.append.return_value = ("OK", []) - # Mock message_exists_in_folder to return False (not a duplicate) - monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False) + # Mock email_exists_in_folder to return False (not a duplicate) + monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False) result = restore_imap_emails.upload_email( mock_conn, @@ -282,19 +282,18 @@ class TestUploadEmail: b"raw email content", '"15-Jan-2024 10:30:00 +0000"', "", - "Test Subject", ) assert result is True mock_conn.append.assert_called_once() def test_upload_email_duplicate(self, monkeypatch): - """Test upload skips duplicate.""" + """Test upload returns False when message exists.""" mock_conn = MagicMock() mock_conn.select.return_value = ("OK", [b"1"]) - # Mock message_exists_in_folder to return True (is a duplicate) - monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: True) + # Mock email_exists_in_folder to return True (is a duplicate) + monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: True) result = restore_imap_emails.upload_email( mock_conn, @@ -302,7 +301,7 @@ class TestUploadEmail: b"raw email content", '"15-Jan-2024 10:30:00 +0000"', "", - "Test Subject", + check_duplicate=True, ) assert result is False @@ -315,8 +314,8 @@ class TestUploadEmail: mock_conn.select.return_value = ("OK", [b"1"]) mock_conn.append.return_value = ("OK", []) - # Mock message_exists_in_folder to return False (not a duplicate) - monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False) + # Mock email_exists_in_folder to return False (not a duplicate) + monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False) result = restore_imap_emails.upload_email( mock_conn, @@ -324,7 +323,6 @@ class TestUploadEmail: b"raw email content", '"15-Jan-2024 10:30:00 +0000"', "", - "Test Subject", flags="\\Seen", # Mark as read ) @@ -504,7 +502,7 @@ class TestGmailModeDraftsFallbackRegression: captured = {} - def fake_upload_email(dest, folder_name, raw_content, date_str, message_id, subject, flags=None): + def fake_upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True): captured["folder_name"] = folder_name return True