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

* Add batch Message-ID fetching for faster duplicate detection

* Remove duplicate get_message_ids_in_folder in migrate_imap_emails

* Optimize orphan email deletion by reusing existing UID mappings

* Remove unused get_msg_details function and related test code

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.

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

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 = []

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

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: <test@example.com>\r\n"), b")"]),
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {b"1": "<test@example.com>"}
+        assert set(result.values()) == {"<test@example.com>"}
+
+    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: <msg1@example.com>\r\n"),
+                    b")",
+                    (b"2 (UID 2 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <msg2@example.com>\r\n"),
+                    b")",
+                    (b"3 (UID 3 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <msg3@example.com>\r\n"),
+                    b")",
+                ],
+            ),
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert set(result.values()) == {"<msg1@example.com>", "<msg2@example.com>", "<msg3@example.com>"}
+
+    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: <valid@example.com>\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()) == {"<valid@example.com>"}

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: <msg1@test>\r\n\r\nBody",
-                b"Subject: Email 2\r\nMessage-ID: <msg2@test>\r\n\r\nBody",
-                b"Subject: Email 3\r\nMessage-ID: <msg3@test>\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 "<msg1@test>" in msg_ids
-            assert "<msg2@test>" in msg_ids
-            assert "<msg3@test>" in msg_ids
-            assert len(msg_ids) == 3
-
-            conn.logout()
-        finally:
-            server.shutdown()
-            thread.join(timeout=2)

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@test.com>",
-            "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@test.com>",
-            "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@test.com>",
-            "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