Fix import path for CI build check (#41)

* Add comprehensive tests for IMAP email restoration functionality

- Implement tests for loading labels manifest, including both old and new formats.
- Add tests for extracting flags from the manifest.
- Create tests for parsing .eml files, including edge cases.
- Develop tests for retrieving .eml files from a directory.
- Validate configuration settings and handle missing credentials.
- Integrate tests for restoring emails from a backup, including Gmail mode and deletion of orphaned emails.
- Test progress caching and backup folder discovery.
- Ensure robust handling of email labels and flags during restoration.
- Verify that the restore process correctly interacts with a mock IMAP server.

* Add comprehensive tests for IMAP utilities and compression support

- Introduced tests for `imap_common.py` covering environment variable verification, IMAP connection handling, folder name normalization, MIME header decoding, message details extraction, duplicate detection, filename sanitization, and trash folder detection.
- Added tests for `imap_compress.py` to validate the functionality of the `_CompressedSocket` wrapper, ensuring data compression and decompression works as expected.
- Implemented integration tests to verify that compression is enabled during IMAP connection setup.

.github/copilot-instructions.md

@@ -0,0 +1,8 @@
+After finishing a batch of changes, always run:
+ruff check src/ tools/ test/
+ruff format --check src/ tools/ test/
+python3 -m pytest test
+
+When writing tests, favor input/output results (integration) over specific implementation and number of calls. Avoid patching and mocking as much as possible.
+
+

.github/workflows/ci.yml

@@ -153,5 +153,5 @@ jobs:
 
       - name: Check imports
         run: |
-          cd src && python -c "import imap_common; import migrate_imap_emails; import backup_imap_emails; import restore_imap_emails; import count_imap_emails; import compare_imap_folders"
+          cd src && python -c "import utils.imap_common; import migrate_imap_emails; import backup_imap_emails; import restore_imap_emails; import count_imap_emails; import compare_imap_folders"
           echo "All imports successful"

.github/workflows/publish.yml

@@ -0,0 +1,59 @@
+name: Publish Python 🐍 distribution to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+    
+    - name: Update version to match release
+      run: |
+        VERSION=${GITHUB_REF_NAME#v}
+        echo "Updating version to $VERSION"
+        sed -i "s/^version = .*/version = \"$VERSION\"/" pyproject.toml
+        grep "^version =" pyproject.toml
+        
+    - name: Install pypa/build
+      run: >-
+        python3 -m pip install build --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: Publish Python 🐍 distribution to PyPI
+    if: startsWith(github.ref, 'refs/tags/')  # only publish to PyPI on tag pushes
+    needs:
+    - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/imap-migration-tools
+    permissions:
+      id-token: write  # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

.gitignore

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

README.md

@@ -16,7 +16,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### The Scripts
 
-1. **`migrate_imap_emails.py`** (The Solution)
+1. **`imap_migrate.py`** (The Solution)
    - Migrates emails folder-by-folder.
    - **Multi-threaded**: Uses a thread pool to copy messages in parallel for high speed.
   - **Smart De-duplication**: Checks if a message already exists in the destination (matching Message-ID) and skips it if found.
@@ -25,19 +25,20 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - **Cleanup**: Optionally deletes messages from the source after successful transfer (effectively a "Move" operation).
      - *Improved for Gmail*: Automatically detects "Trash" folders to ensure emails are properly binned rather than just archived.
    - **Sync Mode**: Optionally deletes emails from destination that no longer exist in source (`--dest-delete`).
+   - **Incremental Cache**: Supports `--migrate-cache <path>` to store a local map of processed emails, drastically speeding up subsequent runs by skipping known messages.
    - **Configurable**: Adjustable concurrency and batch sizes to respect server rate limits.
 
-2. **`compare_imap_folders.py`** (The Validator)
+2. **`imap_compare.py`** (The Validator)
    - Connects to both Source and Destination accounts.
    - Prints a side-by-side comparison table of message counts for every folder.
   - Supports comparing IMAP to a local backup folder (`.eml` files) as either the source or destination.
   - Essential for verifying that the migration was successful and that counts match.
 
-3. **`count_imap_emails.py`** (The Investigator)
+3. **`imap_count.py`** (The Investigator)
   - Counts emails in all folders for a single IMAP account (initial assessment / sizing).
   - Also supports counting a local backup folder (`.eml` files) via `--path` (or `BACKUP_LOCAL_PATH`).
 
-4. **`backup_imap_emails.py`** (The Backup)
+4. **`imap_backup.py`** (The Backup)
    - Downloads emails from an IMAP account to a local disk.
    - **Format**: Saves emails as individual `.eml` files (RFC 5322), compatible with Outlook, Thunderbird, and Apple Mail.
    - **Structure**: Replicates the IMAP folder hierarchy locally.
@@ -45,7 +46,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - **Sync Mode**: Optionally deletes local `.eml` files that no longer exist on the server (`--dest-delete`).
    - **Gmail Labels Preservation**: Creates a `labels_manifest.json` file mapping each email's Message-ID to its Gmail labels, enabling proper restoration with labels intact.
 
-5. **`restore_imap_emails.py`** (The Restore)
+5. **`imap_restore.py`** (The Restore)
    - Uploads emails from a local backup to an IMAP server.
    - **Format**: Reads `.eml` files and uploads them preserving original dates.
    - **Structure**: Recreates the folder hierarchy on the destination server.
@@ -66,7 +67,36 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### 2. Installation
 
-Clone this repository or download the script files to your local machine.
+#### Installation via PyPI (Recommended)
+You can install the tools directly from PyPI.
+
+**For macOS Users (and other externally managed environments):**
+Due to recent changes in Python environments on macOS, it is recommended to use `pipx` to install the tools globally without conflicting with system packages.
+
+1. Install `pipx` (if not already installed):
+   ```bash
+   brew install pipx
+   pipx ensurepath
+   ```
+2. Install the tools:
+   ```bash
+   pipx install imap-migration-tools
+   ```
+
+**For Standard Environments:**
+```bash
+pip install imap-migration-tools
+```
+
+Once installed via PyPI, the following commands are globally available in your terminal:
+- `imap-backup`
+- `imap-compare`
+- `imap-count`
+- `imap-migrate`
+- `imap-restore`
+
+#### Installation from Source
+Alternatively, clone this repository or download the script files to your local machine.
 
 #### macOS
 macOS often comes with Python, but it's best to install the latest version.
@@ -128,7 +158,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```bash
-   python3 migrate_imap_emails.py
+   python3 imap_migrate.py
    ```
 
 #### Windows (PowerShell)
@@ -145,7 +175,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```powershell
-   python migrate_imap_emails.py
+   python imap_migrate.py
    ```
 
 ### Method 2: Command Line Arguments (Overrides)
@@ -153,7 +183,7 @@ All scripts support command-line arguments which take precedence over environmen
 
 **Migration:**
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -166,7 +196,7 @@ python3 migrate_imap_emails.py \
 
 **Comparison:**
 ```bash
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -175,14 +205,14 @@ python3 compare_imap_folders.py \
   --dest-pass "your-app-password"
 
 # Compare IMAP source to a local backup folder
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "./my_backup"
 
 # Compare a local backup folder to an IMAP destination
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -191,23 +221,23 @@ python3 compare_imap_folders.py \
 
 **Counting:**
 ```bash
-python3 count_imap_emails.py \
+python3 imap_count.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --pass "secret"
 
 # Count a local backup folder
-python3 count_imap_emails.py \
+python3 imap_count.py \
    --path "./my_backup"
 
 # Or via environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 count_imap_emails.py
+python3 imap_count.py
 ```
 
 **Counting (OAuth2):**
 ```bash
-python3 count_imap_emails.py \
+python3 imap_count.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --oauth2-client-id "id" \
@@ -218,12 +248,12 @@ export IMAP_HOST="imap.gmail.com"
 export IMAP_USERNAME="me@gmail.com"
 export OAUTH2_CLIENT_ID="id"
 export OAUTH2_CLIENT_SECRET="secret"  # Required for Google
-python3 count_imap_emails.py
+python3 imap_count.py
 ```
 
 **Backup:**
 ```bash
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -235,7 +265,7 @@ python3 backup_imap_emails.py \
 ### 1. Full Migration
 Migrate all folders from Source to Destination.
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -248,7 +278,7 @@ python3 migrate_imap_emails.py \
 For Gmail -> Gmail migrations, `--gmail-mode` migrates only `[Gmail]/All Mail` (no duplicates) and applies labels by copying messages into label folders.
 
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --gmail-mode \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
@@ -258,17 +288,29 @@ python3 migrate_imap_emails.py \
   --dest-pass "dest-app-password"
 ```
 
-Notes:
-- `--preserve-flags` is enabled automatically in `--gmail-mode`.
-- `--dest-delete` is not supported in `--gmail-mode`.
+### 1b. Fast Incremental Migration (Cached)
+Use a local cache file to remember processed emails. Use this for large migrations that may be interrupted or need to run multiple times.
 
-### 1b. Preserve Flags (Any IMAP Server)
+```bash
+python3 imap_migrate.py \
+  --src-host "imap.source.com" \
+  --src-user "source" \
+  --src-pass "pass" \
+  --dest-host "imap.dest.com" \
+  --dest-user "dest" \
+  --dest-pass "pass" \
+  --migrate-cache "./migration_cache"
+```
+
+To force a re-check of cached items without clearing the cache logic entirely, add `--full-migrate`.
+
+### 1c. Preserve Flags (Any IMAP Server)
 Preserve IMAP flags (`\Seen`, `\Flagged`, `\Answered`, `\Draft`) during migration.
 
 If an email already exists on the destination (duplicate), the script can still sync missing flags on the existing message.
 
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --preserve-flags \
   --src-host "imap.example.com" \
   --src-user "source@example.com" \
@@ -281,8 +323,8 @@ python3 migrate_imap_emails.py \
 ### 2. Single Folder Migration
 Migrate ONLY a specific folder (e.g., trying to fix just "Important" or "Sent").
 ```bash
-# Syntax: python3 migrate_imap_emails.py "[Folder Name]"
-python3 migrate_imap_emails.py \
+# Syntax: python3 imap_migrate.py "[Folder Name]"
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -296,7 +338,7 @@ python3 migrate_imap_emails.py \
 Migrate and **delete** from source immediately after verifying the copy.
 ```bash
 # Using flag
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -306,7 +348,7 @@ python3 migrate_imap_emails.py \
   --src-delete
 
 # Or specific folder with delete
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -323,7 +365,7 @@ Keep destination in sync by deleting emails that no longer exist in the source.
 Note: `--dest-delete` is not supported in `--gmail-mode`.
 ```bash
 # Migration: Delete destination emails not found in source
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -333,7 +375,7 @@ python3 migrate_imap_emails.py \
   --dest-delete
 
 # Backup: Delete local .eml files not found on server
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -341,7 +383,7 @@ python3 backup_imap_emails.py \
   --dest-delete
 
 # Restore: Delete server emails not found in local backup
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -354,7 +396,7 @@ python3 restore_imap_emails.py \
 ### 5. Verify Migration
 Compare counts between source and destination.
 ```bash
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -363,14 +405,14 @@ python3 compare_imap_folders.py \
   --dest-pass "dest-app-password"
 
 # IMAP source -> local backup destination
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
   --dest-path "./my_backup"
 
 # local backup source -> IMAP destination
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -388,21 +430,21 @@ INBOX                   | 1250         | 1250       | MATCH
 Download all your emails to your computer as `.eml` files.
 ```bash
 # Backup all folders from an IMAP account to a local folder
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "./backup_folder"
 
 # Or via command line
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "/Users/jdoe/Documents/Emails"
 
 # Backup single folder
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -411,18 +453,18 @@ python3 backup_imap_emails.py \
 ```
 
 ### 6a. Compare IMAP vs Local Backup
-Use `compare_imap_folders.py` to validate an IMAP account against a local backup created by `backup_imap_emails.py`.
+Use `imap_compare.py` to validate an IMAP account against a local backup created by `imap_backup.py`.
 
 ```bash
 # Option 1: IMAP source -> local destination
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "./my_backup"
 
 # Option 2: local source -> IMAP destination
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -437,15 +479,15 @@ export DEST_LOCAL_PATH="./my_backup"
 ```
 
 ### 6b. Count a Local Backup
-Use `count_imap_emails.py` to get per-folder counts from a local backup created by `backup_imap_emails.py`.
+Use `imap_count.py` to get per-folder counts from a local backup created by `imap_backup.py`.
 
 ```bash
 # Option 1: explicit path
-python3 count_imap_emails.py --path "./my_backup"
+python3 imap_count.py --path "./my_backup"
 
 # Option 2: environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 count_imap_emails.py
+python3 imap_count.py
 ```
 
 ### 7. Gmail Backup with Labels Preservation
@@ -453,7 +495,7 @@ When backing up a Gmail account, use `--gmail-mode` for the recommended workflow
 
 ```bash
 # Recommended: Use --gmail-mode for simplest workflow
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -463,7 +505,7 @@ python3 backup_imap_emails.py \
 
 This is equivalent to the more verbose:
 ```bash
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -475,7 +517,7 @@ python3 backup_imap_emails.py \
 **For large accounts (100K+ emails)**, you can build the manifest first to test:
 ```bash
 # Step 1: Build manifest only (fast, no download)
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -483,7 +525,7 @@ python3 backup_imap_emails.py \
   --manifest-only
 
 # Step 2: Download emails (can run later, manifest already exists)
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -529,7 +571,7 @@ python3 backup_imap_emails.py \
 For non-Gmail servers, you can preserve read/starred status with `--preserve-flags`:
 
 ```bash
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.example.com" \
   --src-user "you@example.com" \
   --src-pass "your-password" \
@@ -548,14 +590,14 @@ Use `--full-restore` to force the legacy behavior (process all emails and re-syn
 
 ```bash
 # Restore all folders from backup
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
   --dest-pass "your-app-password"
 
 # Force full restore (legacy behavior)
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -563,7 +605,7 @@ python3 restore_imap_emails.py \
   --full-restore
 
 # Restore with flags (read/starred status)
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.example.com" \
   --dest-user "you@example.com" \
@@ -571,7 +613,7 @@ python3 restore_imap_emails.py \
   --apply-flags
 
 # Restore a specific folder
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -583,7 +625,7 @@ python3 restore_imap_emails.py \
 Restore a Gmail backup with full label structure using `--gmail-mode`:
 
 ```bash
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -600,7 +642,7 @@ python3 restore_imap_emails.py \
 
 **Alternatively**, restore folders individually with labels and flags applied:
 ```bash
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -622,7 +664,7 @@ To use OAuth2, pass `--oauth2-client-id` (or `--src-oauth2-client-id`/`--dest-oa
 
 Environment variable equivalents:
 - Dual-account scripts: `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET` and `DEST_OAUTH2_CLIENT_ID`, `DEST_OAUTH2_CLIENT_SECRET`.
-- Single-account scripts (like `count_imap_emails.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET`).
+- Single-account scripts (like `imap_count.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET`).
 
 ### Microsoft (Outlook / Office 365)
 
@@ -661,7 +703,7 @@ Requires the `msal` package (`pip install msal`). Uses the **device code flow**
 pip install msal
 
 # Migration with Microsoft OAuth2 on source
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "outlook.office365.com" \
   --src-user "user@contoso.com" \
   --src-oauth2-client-id "your-azure-app-client-id" \
@@ -681,7 +723,7 @@ Requires the `google-auth-oauthlib` package (`pip install google-auth-oauthlib`)
 pip install google-auth-oauthlib
 
 # Backup with Google OAuth2
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-oauth2-client-id "your-google-client-id" \
@@ -694,7 +736,7 @@ The script will open your default browser for Google sign-in. After authorizing,
 ## Troubleshooting
 
 - **"Too many simultaneous connections"**:
-  IMAP servers (especially Gmail) limit the number of active connections per IP or user (typically ~15). Since `migrate_imap_emails.py` uses multiple threads, you may hit this limit.
+  IMAP servers (especially Gmail) limit the number of active connections per IP or user (typically ~15). Since `imap_migrate.py` uses multiple threads, you may hit this limit.
   **Solution**: Reduce `MAX_WORKERS` to `4` or `2` using the environment variable.
 
 - **Authentication Errors**:
@@ -747,7 +789,7 @@ PYTHONPATH=src pytest test/ -v
 make coverage
 
 # Run a specific test file
-PYTHONPATH=src pytest test/test_migrate_imap_emails.py -v
+PYTHONPATH=src pytest test/test_imap_migrate.py -v
 
 # Run a specific test
 PYTHONPATH=src pytest test/test_imap_common.py::TestNormalizeFolderName -v
@@ -779,11 +821,11 @@ make ci
 
 | Test File | Description |
 |-----------|-------------|
-| `test_migrate_imap_emails.py` | Email migration tests (basic, duplicates, deletion, folders) |
-| `test_backup_imap_emails.py` | Backup functionality tests |
-| `test_restore_imap_emails.py` | Restore functionality tests |
-| `test_count_imap_emails.py` | Email counting tests |
-| `test_compare_imap_folders.py` | Folder comparison tests |
+| `test_imap_migrate.py` | Email migration tests (basic, duplicates, deletion, folders) |
+| `test_imap_backup.py` | Backup functionality tests |
+| `test_imap_restore.py` | Restore functionality tests |
+| `test_imap_count.py` | Email counting tests |
+| `test_imap_compare.py` | Folder comparison tests |
 | `test_imap_common.py` | Shared utility function tests |
 
 ### Continuous Integration

pyproject.toml

@@ -1,6 +1,44 @@
 # Ruff configuration for IMAP Migration Tools
 # https://docs.astral.sh/ruff/
 
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "imap-migration-tools"
+version = "1.1.0"
+authors = [
+  { name="Javier Callico", email="jcallico@callicode.com" },
+  { name="Nathan Moinvaziri", email="nathan@nathanm.com" }
+]
+description = "Tools for migrating IMAP emails"
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "google-auth-oauthlib",
+    "msal",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/jcallico/imap-migration-tools"
+"Bug Tracker" = "https://github.com/jcallico/imap-migration-tools/issues"
+
+[project.scripts]
+imap-backup = "imap_backup:main"
+imap-restore = "imap_restore:main"
+imap-migrate = "imap_migrate:main"
+imap-count = "imap_count:main"
+imap-compare = "imap_compare:main"
+
+[tool.setuptools.package-dir]
+"" = "src"
+
 [tool.ruff]
 target-version = "py39"
 line-length = 120

src/auth/imap_oauth2.py

@@ -0,0 +1,201 @@
+"""
+IMAP OAuth2 Authentication
+
+OAuth2 token acquisition and refresh for Microsoft and Google IMAP providers.
+Dispatches to provider-specific modules (oauth2_microsoft, oauth2_google).
+
+Provider is auto-detected from the IMAP host string.
+"""
+
+import sys
+import threading
+
+from auth import oauth2_google, oauth2_microsoft
+
+# Re-export caches for test access
+_msal_app_cache = oauth2_microsoft._msal_app_cache
+_google_creds_cache = oauth2_google._creds_cache
+_tenant_cache = oauth2_microsoft._tenant_cache
+
+# Thread-safe token refresh
+_token_refresh_lock = threading.Lock()
+
+
+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.
+    Returns "microsoft", "google", or None if unrecognized.
+    """
+    host_lower = host.lower()
+    if "outlook" in host_lower or "office365" in host_lower or "microsoft" in host_lower:
+        return "microsoft"
+    if "gmail" in host_lower or "google" in host_lower:
+        return "google"
+    return None
+
+
+def discover_microsoft_tenant(email):
+    """
+    Auto-discovers the Microsoft tenant ID from an email address domain.
+    Delegates to oauth2_microsoft module.
+    """
+    return oauth2_microsoft.discover_tenant(email)
+
+
+def acquire_microsoft_oauth2_token(client_id, email):
+    """
+    Acquires a Microsoft OAuth2 access token using the MSAL device code flow.
+    Delegates to oauth2_microsoft module.
+    """
+    return oauth2_microsoft.acquire_token(client_id, email)
+
+
+def acquire_google_oauth2_token(client_id, client_secret):
+    """
+    Acquires a Google OAuth2 access token using the installed app flow.
+    Delegates to oauth2_google module.
+    """
+    return oauth2_google.acquire_token(client_id, client_secret)
+
+
+def acquire_oauth2_token_for_provider(provider, client_id, email, client_secret=None):
+    """
+    Acquires an OAuth2 token for the specified provider.
+
+    Args:
+        provider: "microsoft" or "google"
+        client_id: OAuth2 client ID
+        email: User's email address (used for Microsoft tenant discovery)
+        client_secret: Required for Google, not needed for Microsoft
+    """
+    if provider == "microsoft":
+        return oauth2_microsoft.acquire_token(client_id, email)
+    elif provider == "google":
+        if not client_secret:
+            print(
+                "Error: OAuth2 client secret is required for Google OAuth2. "
+                "Provide --oauth2-client-secret (or --src-oauth2-client-secret / --dest-oauth2-client-secret), "
+                "or set OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET / DEST_OAUTH2_CLIENT_SECRET."
+            )
+            return None
+        return oauth2_google.acquire_token(client_id, client_secret)
+    else:
+        print(f"Error: Unknown OAuth2 provider: {provider}")
+        return None
+
+
+def acquire_token(host, client_id, email, client_secret=None, label=None):
+    """
+    Detect the OAuth2 provider from the host and acquire a token.
+
+    Prints status messages and calls sys.exit(1) on failure.
+
+    Args:
+        host: IMAP host string (used to detect provider)
+        client_id: OAuth2 client ID
+        email: User's email address
+        client_secret: OAuth2 client secret (required for Google)
+        label: Optional context label for messages (e.g. "source", "destination")
+
+    Returns:
+        (token, provider) tuple on success.
+    """
+    provider = detect_oauth2_provider(host)
+    if not provider:
+        print(f"Error: Could not detect OAuth2 provider from host '{host}'.")
+        sys.exit(1)
+    if label:
+        print(f"Acquiring OAuth2 token for {label} ({provider})...")
+    else:
+        print(f"Acquiring OAuth2 token ({provider})...")
+    token = acquire_oauth2_token_for_provider(provider, client_id, email, client_secret)
+    if not token:
+        if label:
+            print(f"Error: Failed to acquire OAuth2 token for {label}.")
+        else:
+            print("Error: Failed to acquire OAuth2 token.")
+        sys.exit(1)
+    if label:
+        print(f"{label.capitalize()} OAuth2 token acquired successfully.\n")
+    else:
+        print("OAuth2 token acquired successfully.\n")
+    return token, provider
+
+
+def auth_description(provider):
+    """
+    Return a human-readable auth description for config summaries.
+
+    Args:
+        provider: OAuth2 provider string (e.g. "microsoft", "google") or None for password auth.
+
+    Returns:
+        "OAuth2/{provider} (XOAUTH2)" or "Basic (password)".
+    """
+    if provider:
+        return f"OAuth2/{provider} (XOAUTH2)"
+    return "Basic (password)"
+
+
+def refresh_oauth2_token(conf, old_token):
+    """
+    Thread-safe OAuth2 token refresh using double-checked locking.
+
+    Multiple threads may detect an expired token simultaneously. This function
+    ensures only one thread performs the actual refresh. Other threads waiting
+    on the lock will see that conf["oauth2_token"] has already been updated and
+    skip the redundant refresh.
+
+    Args:
+        conf: Mutable dict with keys:
+            - host, user, password, oauth2_token: connection credentials
+            - oauth2: dict with provider, client_id, email, client_secret
+        old_token: The expired token that triggered this refresh (for comparison)
+
+    Returns:
+        The new token string, or None if refresh failed.
+    """
+    oauth2 = conf.get("oauth2")
+    if not oauth2:
+        return None
+
+    with _token_refresh_lock:
+        # Another thread may have already refreshed while we were waiting
+        if conf["oauth2_token"] != old_token:
+            return conf["oauth2_token"]
+
+        new_token = acquire_oauth2_token_for_provider(
+            oauth2["provider"], oauth2["client_id"], oauth2["email"], oauth2.get("client_secret")
+        )
+        if new_token:
+            conf["oauth2_token"] = new_token
+        return new_token

src/auth/oauth2_google.py

@@ -0,0 +1,72 @@
+"""
+Google OAuth2 Token Acquisition
+
+OAuth2 token acquisition for Google/Gmail IMAP using installed app flow.
+Opens a browser for user consent and runs a local HTTP server for the redirect.
+
+Requires the 'google-auth-oauthlib' package: pip install google-auth-oauthlib
+"""
+
+import os
+import sys
+
+# Module-level cache for credentials (holds refresh token)
+_creds_cache = {}  # (client_id, client_secret) -> credentials
+
+
+def acquire_token(client_id, client_secret):
+    """
+    Acquires a Google OAuth2 access token using the installed app flow.
+    Opens a browser for user consent and runs a local HTTP server for the redirect.
+    Requires the 'google-auth-oauthlib' package: pip install google-auth-oauthlib
+
+    On subsequent calls, silently refreshes the token using the cached credentials
+    object (which holds the refresh token). No browser interaction needed for refresh.
+    """
+    # Try refreshing cached credentials first (no browser needed)
+    cache_key = (client_id, client_secret)
+    if cache_key in _creds_cache:
+        creds = _creds_cache[cache_key]
+        if creds and creds.refresh_token:
+            try:
+                import google.auth.transport.requests
+
+                creds.refresh(google.auth.transport.requests.Request())
+                if creds.token:
+                    return creds.token
+            except Exception:
+                pass  # Fall through to full auth flow
+
+    try:
+        from google_auth_oauthlib.flow import InstalledAppFlow
+    except ImportError:
+        print("Error: 'google-auth-oauthlib' package is required for Google OAuth2.")
+        print("Install it with: pip install google-auth-oauthlib")
+        sys.exit(1)
+
+    auth_uri = os.getenv("OAUTH2_GOOGLE_AUTH_URL") or "https://accounts.google.com/o/oauth2/auth"
+    token_uri = os.getenv("OAUTH2_GOOGLE_TOKEN_URL") or "https://oauth2.googleapis.com/token"
+
+    client_config = {
+        "installed": {
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "auth_uri": auth_uri,
+            "token_uri": token_uri,
+            "redirect_uris": ["http://localhost"],
+        }
+    }
+
+    flow = InstalledAppFlow.from_client_config(client_config, scopes=["https://mail.google.com/"])
+
+    print("Opening browser for Google authentication...")
+    print("If the browser does not open, check the terminal for a URL to visit.")
+
+    credentials = flow.run_local_server(port=0)
+
+    if credentials and credentials.token:
+        _creds_cache[cache_key] = credentials
+        return credentials.token
+
+    print("Error: Could not acquire Google OAuth2 token.")
+    return None

src/auth/oauth2_microsoft.py

@@ -0,0 +1,154 @@
+"""
+Microsoft OAuth2 Token Acquisition
+
+OAuth2 token acquisition for Microsoft/Outlook IMAP using MSAL device code flow.
+Supports auto-discovery of tenant ID from email domain.
+
+Requires the 'msal' package: pip install msal
+"""
+
+import http.client
+import json
+import os
+import re
+import ssl
+import sys
+import urllib.parse
+
+# Module-level caches
+_msal_app_cache = {}  # (client_id, tenant_id) -> PublicClientApplication
+_tenant_cache = {}  # domain -> tenant_id
+
+
+def _fetch_json_https(host, path, timeout=10):
+    """Fetch JSON from an HTTP(S) endpoint."""
+    if not host or any(ch in host for ch in "\r\n"):
+        raise ValueError("Invalid host")
+    if not path.startswith("/"):
+        path = f"/{path}"
+
+    if host.startswith("http://") or host.startswith("https://"):
+        parsed = urllib.parse.urlparse(host)
+        if not parsed.hostname:
+            raise ValueError("Invalid host")
+        host = parsed.hostname
+        if parsed.port:
+            host = f"{host}:{parsed.port}"
+        base_path = parsed.path.rstrip("/")
+        if base_path:
+            path = f"{base_path}{path}"
+
+        use_https = parsed.scheme == "https"
+    else:
+        use_https = True
+
+    if use_https:
+        context = ssl.create_default_context()
+        conn = http.client.HTTPSConnection(host, timeout=timeout, context=context)
+    else:
+        conn = http.client.HTTPConnection(host, timeout=timeout)
+    try:
+        conn.request("GET", path, headers={"Accept": "application/json"})
+        response = conn.getresponse()
+        body = response.read()
+    finally:
+        conn.close()
+
+    if response.status != 200:
+        raise RuntimeError(f"Unexpected HTTP status {response.status}")
+    return json.loads(body.decode("utf-8"))
+
+
+def discover_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().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"
+
+    discovery_host = os.getenv("OAUTH2_MICROSOFT_DISCOVERY_URL") or "login.microsoftonline.com"
+    try:
+        data = _fetch_json_https(discovery_host, path, timeout=10)
+    except (OSError, http.client.HTTPException, RuntimeError, ValueError) as e:
+        print(f"Error: Could not discover Microsoft tenant for domain '{domain}': {e}")
+        return None
+
+    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:
+        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
+
+
+def acquire_token(client_id, email):
+    """
+    Acquires a Microsoft OAuth2 access token using the MSAL device code flow.
+    Auto-discovers tenant ID from the email domain.
+    Requires the 'msal' package: pip install msal
+
+    On subsequent calls, silently refreshes the token using the cached MSAL app
+    (which holds the refresh token in its in-memory cache).
+    """
+    tenant_id = discover_tenant(email)
+    if not tenant_id:
+        return None
+
+    try:
+        import msal
+    except ImportError:
+        print("Error: 'msal' package is required for Microsoft OAuth2. Install it with: pip install msal")
+        sys.exit(1)
+
+    authority_base = os.getenv("OAUTH2_MICROSOFT_AUTHORITY_BASE_URL")
+    if authority_base:
+        authority = f"{authority_base.rstrip('/')}/{tenant_id}"
+    else:
+        authority = f"https://login.microsoftonline.com/{tenant_id}"
+    scopes = ["https://outlook.office365.com/IMAP.AccessAsUser.All"]
+
+    # Reuse cached MSAL app so acquire_token_silent can access refresh tokens
+    cache_key = (client_id, tenant_id)
+    if cache_key in _msal_app_cache:
+        app = _msal_app_cache[cache_key]
+    else:
+        print(f"Discovered Microsoft tenant: {tenant_id}")
+        app = msal.PublicClientApplication(client_id, authority=authority)
+        _msal_app_cache[cache_key] = app
+
+    # Try cached/refreshed token first (handles refresh tokens automatically)
+    accounts = app.get_accounts()
+    if accounts:
+        result = app.acquire_token_silent(scopes, account=accounts[0])
+        if result and "access_token" in result:
+            return result["access_token"]
+
+    # Fall back to device code flow (first call or if refresh fails)
+    flow = app.initiate_device_flow(scopes=scopes)
+    if "user_code" not in flow:
+        print(f"Error: Could not initiate device flow: {flow.get('error_description', 'Unknown error')}")
+        return None
+
+    print(flow["message"])
+    result = app.acquire_token_by_device_flow(flow)
+
+    if "access_token" in result:
+        return result["access_token"]
+
+    print(f"Error: Could not acquire token: {result.get('error_description', 'Unknown error')}")
+    return None

src/backup_imap_emails.py

@@ -1,894 +1,24 @@
+#!/usr/bin/env python3
 """
-IMAP Email Backup Script
-
-Backs up emails from an IMAP account to a local directory.
-Stores each email as a separate .eml file (RFC 5322 format) which is compatible with
-most email clients (Thunderbird, Apple Mail, Outlook, etc.).
-
-Features:
-- Incremental Backup: Skips messages that have already been downloaded (checks existing UIDs locally).
-- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
-- Folder Replication: Recreates the IMAP folder structure locally.
-- Parallel Processing: Uses multithreading for fast downloads.
-- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
-
-Configuration (Environment Variables):
-    SRC_IMAP_HOST, SRC_IMAP_USERNAME: Source credentials.
-    SRC_IMAP_PASSWORD: Source password (or App Password).
-
-    OAuth2 (Optional - instead of password):
-    SRC_OAUTH2_CLIENT_ID: OAuth2 Client ID
-    SRC_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
-
-  BACKUP_LOCAL_PATH: Destination local directory.
-  MAX_WORKERS: Number of concurrent threads (default: 10).
-  BATCH_SIZE: Number of emails to process per batch (default: 10).
-  PRESERVE_LABELS: Set to "true" to create labels_manifest.json (Gmail). Default is "false".
-  PRESERVE_FLAGS: Set to "true" to preserve IMAP flags in manifest. Default is "false".
-  MANIFEST_ONLY: Set to "true" to only build manifest without downloading. Default is "false".
-  GMAIL_MODE: Set to "true" for Gmail backup mode. Default is "false".
-  DEST_DELETE: Set to "true" to delete local files not found on server (sync mode).
-              Default is "false".
-
-Usage:
-    python3 backup_imap_emails.py \
-        --src-host "imap.example.com" \
-        --src-user "you@example.com" \
-        --src-pass "your-app-password" \
-        --dest-path "./my_backup"
-
-Gmail Labels:
-    python3 backup_imap_emails.py \
-        --src-host "imap.gmail.com" \
-        --src-user "you@gmail.com" \
-        --src-pass "your-app-password" \
-        --dest-path "./my_backup" \
-        --preserve-labels \
-        "[Gmail]/All Mail"
-  This backs up all emails from [Gmail]/All Mail and creates a labels_manifest.json
-  file that maps each email's Message-ID to its Gmail labels for later restoration.
+DEPRECATED: This script has been renamed.
+Please use imap_backup.py instead.
 """
 
-import argparse
-import concurrent.futures
-import json
-import os
 import sys
-import threading
 
-import imap_common
-import imap_oauth2
+from imap_backup import main
 
-# Defaults
-MAX_WORKERS = 10
-BATCH_SIZE = 10
-MANIFEST_FILENAME = "labels_manifest.json"
-
-# Thread-local storage
-thread_local = threading.local()
-print_lock = threading.Lock()
-
-
-def safe_print(message):
-    t_name = threading.current_thread().name
-    short_name = t_name.replace("ThreadPoolExecutor-", "T-").replace("MainThread", "MAIN")
-    with print_lock:
-        print(f"[{short_name}] {message}")
-
-
-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
-    # imap_common.get_imap_connection_from_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]}
-
-    if not hasattr(thread_local, "src") or thread_local.src is None:
-        thread_local.src = imap_common.get_imap_connection_from_conf(src_conf)
+if __name__ == "__main__":
+    print(
+        "WARNING: 'backup_imap_emails.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_backup.py'.",
+        file=sys.stderr,
+    )
+    print("Redirecting to 'imap_backup.py'...", file=sys.stderr)
     try:
-        if thread_local.src:
-            thread_local.src.noop()
-    except Exception:
-        thread_local.src = imap_common.get_imap_connection_from_conf(src_conf)
-        # If reconnection failed (possibly expired token), try refreshing
-        if thread_local.src is None and src_conf.get("oauth2"):
-            old_token = src_conf["oauth2_token"]
-            imap_oauth2.refresh_oauth2_token(src_conf, old_token)
-            thread_local.src = imap_common.get_imap_connection_from_conf(src_conf)
-    return thread_local.src
-
-
-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]}
-
-    src = get_thread_connection(src_conf)
-    if not src:
-        safe_print("Error: Could not establish connection for batch.")
-        return
-
-    try:
-        src.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Error selecting folder {folder_name} in worker: {e}")
-        return
-
-    for uid in uids:
-        try:
-            # Helper to handle byte UIDs
-            uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-
-            # 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
-
-            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)
-                # 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}")
-
-
-def get_existing_uids(local_path):
-    """
-    Scans the local directory for files matching pattern matches {UID}_*.eml
-    Returns a set of UIDs (as strings).
-    """
-    existing = set()
-    if not os.path.exists(local_path):
-        return existing
-
-    try:
-        for filename in os.listdir(local_path):
-            if filename.endswith(".eml") and "_" in filename:
-                # Expecting UID_Subject.eml
-                parts = filename.split("_", 1)
-                if parts[0].isdigit():
-                    existing.add(parts[0])
-    except Exception:
-        pass
-    return existing
-
-
-def is_gmail_label_folder(folder_name):
-    """
-    Determines if a folder represents a Gmail label (user-created or system label
-    that should be preserved).
-    Excludes system folders like All Mail, Spam, Trash, Drafts.
-    """
-    # Exclude system folders that aren't really "labels"
-    if folder_name in imap_common.GMAIL_SYSTEM_FOLDERS:
-        return False
-
-    # imap_common.FOLDER_INBOX is a special case - it's a label in Gmail
-    if folder_name == imap_common.FOLDER_INBOX:
-        return True
-
-    # [Gmail]/Sent Mail and [Gmail]/Starred are labels worth preserving
-    if folder_name in (imap_common.GMAIL_SENT, imap_common.GMAIL_STARRED):
-        return True
-
-    # Any folder NOT under [Gmail]/ is a user label
-    if not folder_name.startswith("[Gmail]/"):
-        return True
-
-    return False
-
-
-# Standard IMAP flags that can be preserved during migration
-# \Recent is session-specific and cannot be set by clients
-# \Deleted should not be preserved as it marks messages for removal
-PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
-
-
-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 = {}
-
-    try:
-        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
-
-    try:
-        resp, data = imap_conn.uid("search", None, "ALL")
-        if resp != "OK" or not data or not data[0]:
-            return message_info
-
-        uids = data[0].split()
-        if not uids:
-            return message_info
-
-        total_uids = len(uids)
-
-        # Fetch Message-IDs and FLAGS in batches - use larger batch for header-only fetches
-        batch_size = 200
-        for i in range(0, len(uids), batch_size):
-            batch = uids[i : i + batch_size]
-            uid_range = b",".join(batch)
-
-            # Report progress
-            if progress_callback:
-                progress_callback(min(i + batch_size, total_uids), total_uids)
-
-            try:
-                # Fetch both Message-ID header and FLAGS
-                resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
-                if resp != "OK":
-                    continue
-
-                # Parse response - items come in pairs for each message
-                for item in items:
-                    if isinstance(item, tuple) and len(item) >= 2:
-                        # First element contains UID and FLAGS info
-                        meta_str = (
-                            item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0])
-                        )
-
-                        # Extract all preservable flags from the metadata
-                        flags = []
-                        for flag in imap_common.PRESERVABLE_FLAGS:
-                            if flag in meta_str:
-                                flags.append(flag)
-
-                        # Second element contains the header
-                        msg_id = imap_common.extract_message_id(item[1])
-                        if msg_id:
-                            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
-
-
-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.
-    This is a convenience wrapper around get_message_info_in_folder.
-    Optional progress_callback(current, total) for progress reporting.
-    """
-    info = get_message_info_in_folder(imap_conn, folder_name, progress_callback)
-    return set(info.keys())
-
-
-def build_labels_manifest(imap_conn, local_path):
-    """
-    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
-    appear in each label, plus their flags from [Gmail]/All Mail.
-
-    Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
-    Saves the manifest to labels_manifest.json in the backup directory.
-    """
-    import time
-
-    manifest = {}
-    start_time = time.time()
-    total_emails_scanned = 0
-
-    safe_print("--- Building Gmail Labels Manifest ---")
-
-    # Get all selectable folders and filter to label folders
-    all_folders = imap_common.list_selectable_folders(imap_conn)
-    if not all_folders:
-        safe_print("Error: Could not list folders for label mapping.")
-        return manifest
-
-    # First, scan [Gmail]/All Mail to get the authoritative flags for all emails
-    safe_print("[0/N] Scanning [Gmail]/All Mail for flags (read/starred/etc)...")
-    all_mail_start = time.time()
-
-    def all_mail_progress_cb(current, total):
-        elapsed = time.time() - start_time
-        print(
-            f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-            end="",
-            flush=True,
-        )
-
-    all_mail_info = get_message_info_in_folder(imap_conn, imap_common.GMAIL_ALL_MAIL, all_mail_progress_cb)
-    print()  # New line after progress
-
-    # Initialize manifest with flags from All Mail
-    for msg_id, info in all_mail_info.items():
-        manifest[msg_id] = {"labels": [], "flags": info.get("flags", [])}
-
-    all_mail_elapsed = time.time() - all_mail_start
-    # Count flag statistics
-    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
-    safe_print(f"  -> {len(all_mail_info)} emails scanned ({all_mail_elapsed:.1f}s)")
-    safe_print(f"  -> Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}\n")
-
-    # Keep connection alive
-    try:
-        imap_conn.noop()
-    except Exception:
-        pass
-
-    # Parse folder names and filter to label folders
-    label_folders = [f for f in all_folders if is_gmail_label_folder(f)]
-
-    total_folders = len(label_folders)
-    safe_print(f"Found {total_folders} label folders to scan.\n")
-
-    # Scan each label folder
-    for folder_idx, folder_name in enumerate(label_folders, 1):
-        folder_start = time.time()
-
-        # Progress callback for this folder
-        def progress_cb(current, total):
-            elapsed = time.time() - start_time
-            print(
-                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-                end="",
-                flush=True,
-            )
-
-        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        message_ids = get_message_ids_in_folder(imap_conn, folder_name, progress_cb)
-        print()  # New line after progress
-
-        # Keep connection alive between folders
-        try:
-            imap_conn.noop()
-        except Exception:
-            pass
-
-        # Determine the label name to store
-        # For [Gmail]/Sent Mail -> "Sent Mail"
-        # For [Gmail]/Starred -> "Starred"
-        # For INBOX -> "INBOX"
-        # For user folders -> folder name as-is
-        if folder_name.startswith("[Gmail]/"):
-            label_name = folder_name[8:]  # Remove "[Gmail]/" prefix
-        else:
-            label_name = folder_name
-
-        for msg_id in message_ids:
-            if msg_id not in manifest:
-                # Email not in All Mail (rare, but handle it)
-                manifest[msg_id] = {"labels": [], "flags": []}
-            if label_name not in manifest[msg_id]["labels"]:
-                manifest[msg_id]["labels"].append(label_name)
-
-        folder_elapsed = time.time() - folder_start
-        total_emails_scanned += len(message_ids)
-        safe_print(f"  -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
-
-    # Summary
-    total_elapsed = time.time() - start_time
-    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
-    safe_print("\nManifest building complete:")
-    safe_print(f"  - Folders scanned: {total_folders + 1}")  # +1 for All Mail
-    safe_print(f"  - Total email-label mappings: {total_emails_scanned}")
-    safe_print(f"  - Unique emails: {len(manifest)}")
-    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}")
-    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
-
-    # Save manifest
-    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
-    try:
-        with open(manifest_path, "w", encoding="utf-8") as f:
-            json.dump(manifest, f, indent=2, ensure_ascii=False)
-        safe_print(f"\nLabels manifest saved to: {manifest_path}")
-    except Exception as e:
-        safe_print(f"Error saving manifest: {e}")
-
-    return manifest
-
-
-def build_flags_manifest(imap_conn, local_path, folders_to_scan=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.
-    """
-    import time
-
-    manifest = {}
-    start_time = time.time()
-
-    safe_print("--- Building Flags Manifest ---")
-
-    # Get folders to scan
-    if folders_to_scan:
-        all_folders = folders_to_scan
-    else:
-        try:
-            typ, folders = imap_conn.list()
-            if typ != "OK":
-                safe_print("Error: Could not list folders.")
-                return manifest
-            all_folders = [imap_common.normalize_folder_name(f) for f in folders]
-        except Exception as e:
-            safe_print(f"Error listing folders: {e}")
-            return manifest
-
-    total_folders = len(all_folders)
-    safe_print(f"Found {total_folders} folders to scan.\n")
-
-    # Scan each folder
-    for folder_idx, folder_name in enumerate(all_folders, 1):
-        folder_start = time.time()
-
-        def progress_cb(current, total):
-            elapsed = time.time() - start_time
-            print(
-                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-                end="",
-                flush=True,
-            )
-
-        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        folder_info = get_message_info_in_folder(imap_conn, folder_name, progress_cb)
-        print()  # New line after progress
-
-        # Merge into manifest (keep first occurrence of flags)
-        for msg_id, info in folder_info.items():
-            if msg_id not in manifest:
-                manifest[msg_id] = {"flags": info.get("flags", [])}
-
-        folder_elapsed = time.time() - folder_start
-        safe_print(f"  -> {len(folder_info)} emails ({folder_elapsed:.1f}s)")
-
-        # Keep connection alive
-        try:
-            imap_conn.noop()
-        except Exception:
-            pass
-
-    # Summary
-    total_elapsed = time.time() - start_time
-    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
-    safe_print("\nFlags manifest building complete:")
-    safe_print(f"  - Folders scanned: {total_folders}")
-    safe_print(f"  - Unique emails: {len(manifest)}")
-    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Flagged: {flagged_count}")
-    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
-
-    # Save manifest
-    manifest_path = os.path.join(local_path, "flags_manifest.json")
-    try:
-        with open(manifest_path, "w", encoding="utf-8") as f:
-            json.dump(manifest, f, indent=2, ensure_ascii=False)
-        safe_print(f"\nFlags manifest saved to: {manifest_path}")
-    except Exception as e:
-        safe_print(f"Error saving manifest: {e}")
-
-    return manifest
-
-
-def load_labels_manifest(local_path):
-    """
-    Loads an existing labels manifest from the backup directory.
-    Returns the manifest dict or empty dict if not found.
-    """
-    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
-    if os.path.exists(manifest_path):
-        try:
-            with open(manifest_path, encoding="utf-8") as f:
-                return json.load(f)
-        except Exception:
-            pass
-    return {}
-
-
-def delete_orphan_local_files(local_folder_path, server_uids):
-    """
-    Delete local .eml files that no longer exist on the server.
-    Args:
-        local_folder_path: Path to local folder containing .eml files
-        server_uids: Set of UID strings currently on server
-    Returns:
-        Count of deleted files
-    """
-    deleted_count = 0
-    if not os.path.exists(local_folder_path):
-        return deleted_count
-
-    try:
-        for filename in os.listdir(local_folder_path):
-            if not filename.endswith(".eml") or "_" not in filename:
-                continue
-
-            # Extract UID from filename (format: {UID}_{Subject}.eml)
-            parts = filename.split("_", 1)
-            if not parts[0].isdigit():
-                continue
-
-            local_uid = parts[0]
-            if local_uid not in server_uids:
-                file_path = os.path.join(local_folder_path, filename)
-                try:
-                    os.remove(file_path)
-                    safe_print(f"  -> Deleted orphan: {filename}")
-                    deleted_count += 1
-                except Exception as e:
-                    safe_print(f"  -> Error deleting {filename}: {e}")
-    except Exception as e:
-        safe_print(f"Error scanning for orphan files: {e}")
-
-    return deleted_count
-
-
-def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=False):
-    safe_print(f"--- Processing Folder: {folder_name} ---")
-
-    # create local path
-    # Handle folder separators. IMAP output might be "Parent/Child"
-    # We rely on OS to handle "Parent/Child" as subdirectories using join
-    # But clean the segments
-    cleaned_name = folder_name.replace("/", os.sep)
-    local_folder_path = os.path.join(local_base_path, cleaned_name)
-
-    try:
-        os.makedirs(local_folder_path, exist_ok=True)
-    except Exception as e:
-        safe_print(f"Error creating directory {local_folder_path}: {e}")
-        return
-
-    # Select IMAP folder
-    try:
-        src_main.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Skipping {folder_name}: {e}")
-        return
-
-    # Search all
-    resp, data = src_main.uid("search", None, "ALL")
-    if resp != "OK":
-        return
-
-    uids = data[0].split()
-    total_on_server = len(uids)
-
-    # Build set of server UIDs for comparison
-    server_uid_set = set()
-    for u in uids:
-        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
-        server_uid_set.add(u_str)
-
-    if total_on_server == 0:
-        safe_print(f"Folder {folder_name} is empty.")
-        # If dest_delete enabled, delete all local files
-        if dest_delete:
-            deleted = delete_orphan_local_files(local_folder_path, set())
-            if deleted > 0:
-                safe_print(f"Deleted {deleted} orphan files from local backup.")
-        return
-
-    # Incremental Optimization
-    # Read local directory to find UIDs we already have
-    existing_uids = get_existing_uids(local_folder_path)
-
-    # Delete orphan local files if dest_delete is enabled
-    if dest_delete:
-        orphan_uids = existing_uids - server_uid_set
-        if orphan_uids:
-            safe_print(f"Found {len(orphan_uids)} local files not on server, deleting...")
-            deleted = delete_orphan_local_files(local_folder_path, server_uid_set)
-            if deleted > 0:
-                safe_print(f"Deleted {deleted} orphan files from local backup.")
-
-    # Filter UIDs
-    # decode uid first if bytes
-    uids_to_download = []
-    for u in uids:
-        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
-        if u_str not in existing_uids:
-            uids_to_download.append(u)
-
-    skipped = total_on_server - len(uids_to_download)
-    if skipped > 0:
-        safe_print(f"Skipping {skipped} emails (already exist locally).")
-
-    if not uids_to_download:
-        safe_print(f"Folder {folder_name} is up to date.")
-        return
-
-    safe_print(f"Downloading {len(uids_to_download)} new emails...")
-
-    # Create batches
-    uid_batches = [uids_to_download[i : i + BATCH_SIZE] for i in range(0, len(uids_to_download), BATCH_SIZE)]
-
-    executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
-    try:
-        futures = []
-        for batch in uid_batches:
-            futures.append(executor.submit(process_batch, batch, folder_name, src_conf, local_folder_path))
-
-        for future in concurrent.futures.as_completed(futures):
-            future.result()
-
+        main()
     except KeyboardInterrupt:
-        executor.shutdown(wait=False, cancel_futures=True)
-        raise
-    finally:
-        executor.shutdown(wait=True)
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
-
-    # Source
-    default_src_host = os.getenv("SRC_IMAP_HOST")
-    default_src_user = os.getenv("SRC_IMAP_USERNAME")
-    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
-    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--src-host",
-        default=default_src_host,
-        required=not bool(default_src_host),
-        help="Source IMAP Server (or SRC_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--src-user",
-        default=default_src_user,
-        required=not bool(default_src_user),
-        help="Source Username (or SRC_IMAP_USERNAME)",
-    )
-
-    # Authentication: require either password OR OAuth2 client-id (unless provided via env vars)
-    auth_required = not bool(default_src_pass or default_src_client_id)
-    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
-    auth_group.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
-    # OAuth2
-    auth_group.add_argument(
-        "--src-oauth2-client-id",
-        default=default_src_client_id,
-        dest="src_client_id",
-        help="OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
-    )
-    parser.add_argument(
-        "--src-oauth2-client-secret",
-        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="src_client_secret",
-        help="OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
-    )
-
-    # Destination (Local Path)
-    env_path = os.getenv("BACKUP_LOCAL_PATH")
-    parser.add_argument(
-        "--dest-path",
-        default=env_path,
-        required=not bool(env_path),
-        help="Local destination path (or BACKUP_LOCAL_PATH)",
-    )
-
-    # Config
-    parser.add_argument("--workers", type=int, default=int(os.getenv("MAX_WORKERS", 10)), help="Thread count")
-    parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
-
-    # Gmail Labels
-    env_preserve_labels = os.getenv("PRESERVE_LABELS", "false").lower() == "true"
-    parser.add_argument(
-        "--preserve-labels",
-        action="store_true",
-        default=env_preserve_labels,
-        help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
-    )
-    env_preserve_flags = os.getenv("PRESERVE_FLAGS", "false").lower() == "true"
-    parser.add_argument(
-        "--preserve-flags",
-        action="store_true",
-        default=env_preserve_flags,
-        help="Preserve IMAP flags (read/unread, starred, answered, draft) in manifest for restoration",
-    )
-    env_manifest_only = os.getenv("MANIFEST_ONLY", "false").lower() == "true"
-    parser.add_argument(
-        "--manifest-only",
-        action="store_true",
-        default=env_manifest_only,
-        help="Gmail only: Build the labels manifest and exit without downloading emails",
-    )
-    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
-    parser.add_argument(
-        "--gmail-mode",
-        action="store_true",
-        default=env_gmail_mode,
-        help="Gmail backup mode: Build labels manifest and backup [Gmail]/All Mail only (recommended)",
-    )
-
-    # Sync mode: delete local files not on server
-    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
-    parser.add_argument(
-        "--dest-delete",
-        action="store_true",
-        default=env_dest_delete,
-        help="Delete local .eml files that no longer exist on the IMAP server (sync mode)",
-    )
-
-    parser.add_argument("folder", nargs="?", help="Specific folder to backup")
-
-    args = parser.parse_args()
-
-    use_oauth2 = bool(args.src_client_id)
-
-    global MAX_WORKERS, BATCH_SIZE
-    MAX_WORKERS = args.workers
-    BATCH_SIZE = args.batch
-
-    # Acquire OAuth2 token if configured
-    oauth2_token = None
-    oauth2_provider = None
-    if use_oauth2:
-        oauth2_provider = imap_oauth2.detect_oauth2_provider(args.src_host)
-        if not oauth2_provider:
-            print(f"Error: Could not detect OAuth2 provider from host '{args.src_host}'.")
-            sys.exit(1)
-        print(f"Acquiring OAuth2 token ({oauth2_provider})...")
-        oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
-            oauth2_provider, args.src_client_id, args.src_user, args.src_client_secret
-        )
-        if not oauth2_token:
-            print("Error: Failed to acquire OAuth2 token.")
-            sys.exit(1)
-        print("OAuth2 token acquired successfully.\n")
-
-    # Use a dict so token updates propagate to worker threads
-    src_conf = {
-        "host": args.src_host,
-        "user": args.src_user,
-        "password": args.src_pass,
-        "oauth2_token": oauth2_token,
-        "oauth2": {
-            "provider": oauth2_provider,
-            "client_id": args.src_client_id,
-            "email": args.src_user,
-            "client_secret": args.src_client_secret,
-        }
-        if use_oauth2
-        else None,
-    }
-
-    # Expand path (~/...)
-    local_path = os.path.expanduser(args.dest_path)
-    if not os.path.exists(local_path):
-        try:
-            os.makedirs(local_path)
-            print(f"Created backup directory: {local_path}")
-        except Exception as e:
-            print(f"Error creating backup directory: {e}")
-            sys.exit(1)
-
-    print("\n--- Configuration Summary ---")
-    print(f"Source Host     : {args.src_host}")
-    print(f"Source User     : {args.src_user}")
-    print(f"Auth Method     : {'OAuth2/' + oauth2_provider + ' (XOAUTH2)' if use_oauth2 else 'Basic (password)'}")
-    print(f"Destination Path: {local_path}")
-    if args.gmail_mode:
-        print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
-    elif args.manifest_only:
-        print("Mode            : Manifest Only (no email download)")
-    elif args.folder:
-        print(f"Target Folder   : {args.folder}")
-    if args.preserve_labels or args.manifest_only or args.gmail_mode:
-        print("Preserve Labels : Yes (Gmail)")
-    if args.preserve_flags or args.gmail_mode:
-        print("Preserve Flags  : Yes (read/starred/answered/draft)")
-    if args.dest_delete:
-        print("Dest Delete     : Yes (remove local orphans)")
-    print("-----------------------------\n")
-
-    try:
-        src = imap_common.get_imap_connection_from_conf(src_conf)
-        if not src:
-            sys.exit(1)
-
-        # Build labels manifest BEFORE backing up emails (Gmail mode)
-        # This way we capture the label state at backup time
-        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)
-            print("")  # Blank line after manifest building
-        # Build flags-only manifest for non-Gmail servers
-        elif args.preserve_flags:
-            print("Building flags manifest...")
-            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)
-            print("")  # Blank line after manifest building
-
-        # If manifest-only mode, we're done
-        if args.manifest_only:
-            src.logout()
-            manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
-            print("\nManifest-only mode complete.")
-            print(f"Labels manifest saved to: {manifest_path}")
-            print("\nTo download emails, run again without --manifest-only:")
-            print(f'  python3 backup_imap_emails.py --dest-path "{local_path}" "[Gmail]/All Mail"')
-            sys.exit(0)
-
-        # Gmail mode: backup only [Gmail]/All Mail
-        if args.gmail_mode:
-            backup_folder(src, imap_common.GMAIL_ALL_MAIL, local_path, src_conf, args.dest_delete)
-        elif args.folder:
-            backup_folder(src, args.folder, local_path, src_conf, args.dest_delete)
-        else:
-            folders = imap_common.list_selectable_folders(src)
-            for name in folders:
-                try:
-                    src.noop()
-                except Exception:
-                    if src_conf.get("oauth2"):
-                        safe_print("Refreshing OAuth2 token...")
-                        imap_oauth2.refresh_oauth2_token(src_conf, src_conf["oauth2_token"])
-
-                backup_folder(src, name, local_path, src_conf, args.dest_delete)
-
-        src.logout()
-        print("\nBackup completed successfully.")
-
-        if args.preserve_labels or args.gmail_mode:
-            manifest_path = os.path.join(local_path, "labels_manifest.json")
-            print(f"\nGmail labels manifest saved to: {manifest_path}")
-            print("Use this file when restoring to reapply labels and flags to emails.")
-        elif args.preserve_flags:
-            manifest_path = os.path.join(local_path, "flags_manifest.json")
-            print(f"\nFlags manifest saved to: {manifest_path}")
-            print("Use this file when restoring to reapply read/starred status to emails.")
-
-    except KeyboardInterrupt:
-        print("\nBackup interrupted by user.")
+        print("\nProcess terminated by user.")
         sys.exit(0)
     except Exception as e:
         print(f"Fatal Error: {e}")
-
-
-if __name__ == "__main__":
-    main()
+        sys.exit(1)

src/compare_imap_folders.py

@@ -1,406 +1,24 @@
+#!/usr/bin/env python3
 """
-IMAP Folder Comparison Script
-
-This script compares email counts between a source and a destination.
-Each side can be either an IMAP account or a local backup folder.
-It iterates through all folders found in the source account and checks the corresponding
-folder in the destination account.
-
-Configuration (Environment Variables):
-  Source Account:
-    SRC_IMAP_HOST       : Source IMAP Host
-    SRC_IMAP_USERNAME   : Source Username/Email
-    SRC_IMAP_PASSWORD   : Source Password
-
-    OAuth2 (Optional - instead of password):
-    SRC_OAUTH2_CLIENT_ID     : OAuth2 Client ID
-    SRC_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
-
-  Destination Account:
-    DEST_IMAP_HOST      : Destination IMAP Host
-    DEST_IMAP_USERNAME  : Destination Username/Email
-    DEST_IMAP_PASSWORD  : Destination Password
-
-    OAuth2 (Optional - instead of password):
-    DEST_OAUTH2_CLIENT_ID     : OAuth2 Client ID
-    DEST_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
-
-Also supports local folders as source and/or destination:
-    SRC_LOCAL_PATH      : Source local folder (backup root)
-    DEST_LOCAL_PATH     : Destination local folder (backup root)
-
-Usage:
-  python3 compare_imap_folders.py
-
-Examples:
-        # IMAP -> IMAP
-        python3 compare_imap_folders.py \
-            --src-host "imap.source.com" \
-            --src-user "source@example.com" \
-            --src-pass "source-app-password" \
-            --dest-host "imap.dest.com" \
-            --dest-user "dest@example.com" \
-            --dest-pass "dest-app-password"
-
-        # Local -> IMAP
-        python3 compare_imap_folders.py \
-            --src-path "./my_backup" \
-            --dest-host "imap.dest.com" \
-            --dest-user "dest@example.com" \
-            --dest-pass "dest-app-password"
-
-        # IMAP -> Local
-        python3 compare_imap_folders.py \
-            --src-host "imap.source.com" \
-            --src-user "source@example.com" \
-            --src-pass "source-app-password" \
-            --dest-path "./my_backup"
+DEPRECATED: This script has been renamed.
+Please use imap_compare.py instead.
 """
 
-import argparse
-import os
 import sys
-from typing import Optional
-
-import imap_common
-import imap_oauth2
-
-
-def get_email_count(conn, folder_name):
-    """Return the IMAP message count for a folder, or None on error."""
-    try:
-        # Select folder in read-only mode
-        # Quote folder name handles spaces
-        typ, data = conn.select(f'"{folder_name}"', readonly=True)
-        if typ != "OK":
-            return None
-
-        # SELECT command returns the number of messages in data[0]
-        # data[0] is bytes, e.g. b'123'
-        if data and data[0]:
-            return int(data[0])
-        return 0
-
-    except Exception:
-        # print(f"Error checking {folder_name}: {e}")
-        return None
-
-
-def _is_ignored_local_dir(dirname: str) -> bool:
-    return dirname.startswith(".") or dirname == "__pycache__"
-
-
-def list_local_folders(local_root: str) -> list[str]:
-    """List all folders under a local backup root in IMAP-style names.
-
-    The local backup format is expected to mirror IMAP folder hierarchy using
-    subdirectories (e.g. "[Gmail]/All Mail" becomes "[Gmail]/All Mail/").
-    """
-    folders: set[str] = set()
-
-    for dirpath, dirnames, _filenames in os.walk(local_root):
-        dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)]
-
-        if os.path.abspath(dirpath) == os.path.abspath(local_root):
-            continue
-
-        rel = os.path.relpath(dirpath, local_root)
-        if rel == ".":
-            continue
-
-        parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)]
-        if not parts:
-            continue
-
-        folders.add("/".join(parts))
-
-    return sorted(folders)
-
-
-def get_local_email_count(local_root: str, folder_name: str) -> Optional[int]:
-    """Return the count of .eml files in a local folder, or None if missing/unreadable."""
-    folder_path = os.path.join(local_root, *folder_name.split("/"))
-    if not os.path.isdir(folder_path):
-        return None
-
-    try:
-        count = 0
-        for filename in os.listdir(folder_path):
-            if not filename.endswith(".eml"):
-                continue
-            full_path = os.path.join(folder_path, filename)
-            if os.path.isfile(full_path):
-                count += 1
-        return count
-    except OSError:
-        return None
-
-
-def main():
-    default_src_path = os.getenv("SRC_LOCAL_PATH")
-    default_dest_path = os.getenv("DEST_LOCAL_PATH")
-
-    # Phase 1: determine whether each side is local
-    pre_parser = argparse.ArgumentParser(add_help=False)
-    pre_parser.add_argument("--src-path", default=default_src_path)
-    pre_parser.add_argument("--dest-path", default=default_dest_path)
-    pre_args, _ = pre_parser.parse_known_args()
-
-    src_requires_imap = not bool(pre_args.src_path)
-    dest_requires_imap = not bool(pre_args.dest_path)
-
-    # Phase 2: full parser with conditional requirements
-    parser = argparse.ArgumentParser(description="Compare email counts between two IMAP accounts.")
-    parser.add_argument(
-        "--src-path",
-        default=default_src_path,
-        help="Source local folder (backup root). If set, IMAP source args are ignored.",
-    )
-    parser.add_argument(
-        "--dest-path",
-        default=default_dest_path,
-        help="Destination local folder (backup root). If set, IMAP destination args are ignored.",
-    )
-
-    # Source args
-    default_src_host = os.getenv("SRC_IMAP_HOST")
-    default_src_user = os.getenv("SRC_IMAP_USERNAME")
-    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
-    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--src-host",
-        default=default_src_host,
-        required=src_requires_imap and not bool(default_src_host),
-        help="Source IMAP Server (or SRC_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--src-user",
-        default=default_src_user,
-        required=src_requires_imap and not bool(default_src_user),
-        help="Source Username (or SRC_IMAP_USERNAME)",
-    )
-    src_auth_required = src_requires_imap and not bool(default_src_pass or default_src_client_id)
-    src_auth = parser.add_mutually_exclusive_group(required=src_auth_required)
-    src_auth.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
-    src_auth.add_argument(
-        "--src-oauth2-client-id",
-        default=default_src_client_id,
-        dest="src_client_id",
-        help="Source OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
-    )
-    parser.add_argument(
-        "--src-oauth2-client-secret",
-        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="src_client_secret",
-        help="Source OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
-    )
-
-    # Dest args
-    default_dest_host = os.getenv("DEST_IMAP_HOST")
-    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
-    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
-    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--dest-host",
-        default=default_dest_host,
-        required=dest_requires_imap and not bool(default_dest_host),
-        help="Destination IMAP Server (or DEST_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--dest-user",
-        default=default_dest_user,
-        required=dest_requires_imap and not bool(default_dest_user),
-        help="Destination Username (or DEST_IMAP_USERNAME)",
-    )
-    dest_auth_required = dest_requires_imap and not bool(default_dest_pass or default_dest_client_id)
-    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
-    dest_auth.add_argument(
-        "--dest-pass",
-        default=default_dest_pass,
-        help="Destination Password (or DEST_IMAP_PASSWORD)",
-    )
-    dest_auth.add_argument(
-        "--dest-oauth2-client-id",
-        default=default_dest_client_id,
-        dest="dest_client_id",
-        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
-    )
-    dest_auth.add_argument(
-        "--dest-client-id",
-        default=default_dest_client_id,
-        dest="dest_client_id",
-        help=argparse.SUPPRESS,
-    )
-    parser.add_argument(
-        "--dest-oauth2-client-secret",
-        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
-        dest="dest_client_secret",
-        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
-    )
-    parser.add_argument(
-        "--dest-client-secret",
-        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
-        dest="dest_client_secret",
-        help=argparse.SUPPRESS,
-    )
-
-    args = parser.parse_args()
-
-    src_is_local = bool(args.src_path)
-    dest_is_local = bool(args.dest_path)
-
-    SRC_HOST = args.src_host
-    SRC_USER = args.src_user
-    DEST_HOST = args.dest_host
-    DEST_USER = args.dest_user
-
-    src_use_oauth2 = bool(args.src_client_id) and not src_is_local
-    dest_use_oauth2 = bool(args.dest_client_id) and not dest_is_local
-
-    # Acquire OAuth2 tokens if configured
-    src_oauth2_token = None
-    src_oauth2_provider = None
-    if src_use_oauth2:
-        src_oauth2_provider = imap_oauth2.detect_oauth2_provider(SRC_HOST)
-        if not src_oauth2_provider:
-            print(f"Error: Could not detect OAuth2 provider from host '{SRC_HOST}'.")
-            sys.exit(1)
-        print(f"Acquiring OAuth2 token for source ({src_oauth2_provider})...")
-        src_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
-            src_oauth2_provider, args.src_client_id, SRC_USER, args.src_client_secret
-        )
-        if not src_oauth2_token:
-            print("Error: Failed to acquire OAuth2 token for source.")
-            sys.exit(1)
-        print("Source OAuth2 token acquired successfully.\n")
-
-    dest_oauth2_token = None
-    dest_oauth2_provider = None
-    if dest_use_oauth2:
-        dest_oauth2_provider = imap_oauth2.detect_oauth2_provider(DEST_HOST)
-        if not dest_oauth2_provider:
-            print(f"Error: Could not detect OAuth2 provider from host '{DEST_HOST}'.")
-            sys.exit(1)
-        print(f"Acquiring OAuth2 token for destination ({dest_oauth2_provider})...")
-        dest_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
-            dest_oauth2_provider, args.dest_client_id, DEST_USER, args.dest_client_secret
-        )
-        if not dest_oauth2_token:
-            print("Error: Failed to acquire OAuth2 token for destination.")
-            sys.exit(1)
-        print("Destination OAuth2 token acquired successfully.\n")
-
-    print("\n--- Configuration Summary ---")
-    if src_is_local:
-        print(f"Source (Local)  : {args.src_path}")
-    else:
-        print(f"Source Host     : {args.src_host}")
-        print(f"Source User     : {args.src_user}")
-        print(
-            f"Source Auth     : {'OAuth2/' + src_oauth2_provider + ' (XOAUTH2)' if src_use_oauth2 else 'Basic (password)'}"
-        )
-
-    if dest_is_local:
-        print(f"Destination (Local): {args.dest_path}")
-    else:
-        print(f"Destination Host: {args.dest_host}")
-        print(f"Destination User: {args.dest_user}")
-        print(
-            f"Destination Auth: {'OAuth2/' + dest_oauth2_provider + ' (XOAUTH2)' if dest_use_oauth2 else 'Basic (password)'}"
-        )
-    print("-----------------------------\n")
-
-    src = None
-    dest = None
-
-    try:
-        if not src_is_local:
-            # Connect to Source
-            print("Connecting to Source...")
-            src = imap_common.get_imap_connection(args.src_host, args.src_user, args.src_pass, src_oauth2_token)
-            if not src:
-                return
-
-        if not dest_is_local:
-            # Connect to Dest
-            print("Connecting to Destination...")
-            dest = imap_common.get_imap_connection(args.dest_host, args.dest_user, args.dest_pass, dest_oauth2_token)
-            if not dest:
-                return
-
-        # List Source Folders
-        print("Listing folders in Source...")
-        if src_is_local:
-            folders = list_local_folders(args.src_path)
-        else:
-            folders = imap_common.list_selectable_folders(src)
-
-        if not folders:
-            print("Failed to list source folders.")
-            return
-
-        # Prepare Table Header
-        header = f"{'Folder Name':<40} | {'Source':>10} | {'Dest':>10} | {'Diff':>10}"
-        print("-" * len(header))
-        print(header)
-        print("-" * len(header))
-
-        total_src = 0
-        total_dest = 0
-
-        # Iterate through Source folders
-        for folder_name in folders:
-            # Get Counts
-            if src_is_local:
-                src_count = get_local_email_count(args.src_path, folder_name)
-            else:
-                src_count = get_email_count(src, folder_name)
-
-            if dest_is_local:
-                dest_count = get_local_email_count(args.dest_path, folder_name)
-            else:
-                dest_count = get_email_count(dest, folder_name)
-
-            # Format for display
-            src_str = str(src_count) if src_count is not None else "Err"
-            dest_str = str(dest_count) if dest_count is not None else "N/A"  # N/A usually means folder doesn't exist
-
-            diff_str = ""
-            if src_count is not None and dest_count is not None:
-                diff = src_count - dest_count
-                diff_str = str(diff)
-                total_src += src_count
-                total_dest += dest_count
-            elif src_count is not None:
-                total_src += src_count
-
-            print(f"{folder_name:<40} | {src_str:>10} | {dest_str:>10} | {diff_str:>10}")
-
-        print("-" * len(header))
-        print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
-
-    except Exception as e:
-        print(f"\nFatal Error: {e}")
-        if "Too many simultaneous connections" in str(e):
-            print("Tip: Wait a few minutes for previous connections to timeout or check other active scripts.")
-
-    finally:
-        # Check source connection state and logout if possible
-        if src:
-            try:
-                src.logout()
-            except Exception:
-                pass
-
-        # Check dest connection state and logout if possible
-        if dest:
-            try:
-                dest.logout()
-            except Exception:
-                pass
 
+from imap_compare import main
 
 if __name__ == "__main__":
-    main()
+    print(
+        "WARNING: 'compare_imap_folders.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_compare.py'.",
+        file=sys.stderr,
+    )
+    print("Redirecting to 'imap_compare.py'...", file=sys.stderr)
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/core/imap_session.py

@@ -0,0 +1,123 @@
+"""
+IMAP Session Management
+
+Connection and session management for IMAP operations with OAuth2 support.
+Combines imap_common (low-level IMAP) with imap_oauth2 (token refresh).
+"""
+
+from auth import imap_oauth2
+from utils import imap_common
+
+
+def build_imap_conf(host, user, password, client_id=None, client_secret=None, label=None):
+    """
+    Build a standard IMAP connection config dict.
+
+    If client_id is provided, acquires an OAuth2 token (with error handling and
+    sys.exit(1) on failure). Otherwise, builds a password-auth config.
+
+    Args:
+        host: IMAP host
+        user: IMAP username / email
+        password: IMAP password (used for password auth or as fallback)
+        client_id: OAuth2 client ID (triggers OAuth2 flow if provided)
+        client_secret: OAuth2 client secret (required for Google)
+        label: Optional context label for status messages (e.g. "source", "destination")
+
+    Returns:
+        Dict with keys: host, user, password, oauth2_token, oauth2
+    """
+    oauth2_token = None
+    oauth2_provider = None
+    oauth2_info = None
+
+    if client_id:
+        oauth2_token, oauth2_provider = imap_oauth2.acquire_token(host, client_id, user, client_secret, label)
+        oauth2_info = {
+            "provider": oauth2_provider,
+            "client_id": client_id,
+            "email": user,
+            "client_secret": client_secret,
+        }
+
+    return {
+        "host": host,
+        "user": user,
+        "password": password,
+        "oauth2_token": oauth2_token,
+        "oauth2": oauth2_info,
+    }
+
+
+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.
+    """
+    # The OAuth2 provider implementations (MSAL for Microsoft, google-auth for Google)
+    # use internal caching and will only contact the server if the token needs refresh.
+    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
+
+
+def get_thread_connection(thread_store, key, conf):
+    """Get or refresh a thread-local IMAP connection.
+
+    Stores the connection on thread_store under the given key so it persists
+    across calls within the same thread.
+
+    Args:
+        thread_store: A threading.local() instance
+        key: Attribute name to store the connection under (e.g. "src", "dest")
+        conf: Connection config dict (see ensure_connection)
+
+    Returns:
+        Healthy IMAP connection, or None if connection failed.
+    """
+    conn = ensure_connection(getattr(thread_store, key, None), conf)
+    setattr(thread_store, key, conn)
+    return conn

src/count_imap_emails.py

@@ -1,296 +1,24 @@
-"""IMAP Email Counting Script.
-
-Counts emails per folder from either:
-- An IMAP account, or
-- A local backup folder created by ``backup_imap_emails.py`` (counts ``.eml`` files).
-
-Configuration (Environment Variables):
-    IMAP_HOST        : IMAP Host (e.g., imap.gmail.com)
-    IMAP_USERNAME    : Username/Email
-    IMAP_PASSWORD    : Password (or App Password)
-
-    OAuth2 (Optional - instead of password):
-    OAUTH2_CLIENT_ID        : OAuth2 Client ID
-    OAUTH2_CLIENT_SECRET    : OAuth2 Client Secret (required for Google)
-    SRC_OAUTH2_CLIENT_ID    : Alternate OAuth2 client ID env var
-    SRC_OAUTH2_CLIENT_SECRET: Alternate OAuth2 client secret env var
-
-Local backup counting:
-    BACKUP_LOCAL_PATH : Local backup root (preferred)
-    SRC_LOCAL_PATH    : Alternate local backup root
-
-Examples:
-    # Count an IMAP account
-    export IMAP_HOST="imap.gmail.com"
-    export IMAP_USERNAME="user@gmail.com"
-    export IMAP_PASSWORD="secretpassword"
-    python3 count_imap_emails.py
-
-    # Count an IMAP account using OAuth2
-    export IMAP_HOST="imap.gmail.com"
-    export IMAP_USERNAME="user@gmail.com"
-    export OAUTH2_CLIENT_ID="your-client-id"
-    export OAUTH2_CLIENT_SECRET="your-client-secret"  # Required for Google
-    python3 count_imap_emails.py
-
-    # Count a local backup
-    python3 count_imap_emails.py --path "./my_backup"
-
-    # Or set a default local backup path via env var
-    export BACKUP_LOCAL_PATH="./my_backup"
-    python3 count_imap_emails.py
+#!/usr/bin/env python3
+"""
+DEPRECATED: This script has been renamed.
+Please use imap_count.py instead.
 """
 
-import argparse
-import imaplib
-import os
 import sys
-from typing import Optional
-
-import imap_common
-import imap_oauth2
-
-
-def count_emails(imap_server, username, password=None, oauth2_token=None):
-    try:
-        # Connect to the IMAP server (using SSL)
-        print(f"Connecting to {imap_server}...")
-        mail = imap_common.get_imap_connection(imap_server, username, password, oauth2_token)
-        if not mail:
-            return
-
-        # List all mailboxes
-        print("Listing mailboxes...")
-        folders = imap_common.list_selectable_folders(mail)
-
-        if not folders:
-            print("Failed to list mailboxes.")
-            return
-
-        total_all_folders = 0
-        print(f"{'Folder Name':<40} {'Count':>10}")
-        print("-" * 52)
-
-        for folder_name in folders:
-            display_name = folder_name
-
-            try:
-                # Select the mailbox (read-only is sufficient for counting)
-                # folder_name extracted from list usually handles quotes correctly for select
-                rv, _ = mail.select(f'"{folder_name}"', readonly=True)
-                if rv != "OK":
-                    print(f"{display_name:<40} {'Skipped':>10}")
-                    continue
-
-                # Search for all emails
-                status, data = mail.search(None, "ALL")
-
-                if status == "OK":
-                    # data[0] is space separated IDs
-                    email_ids = data[0].split()
-                    count = len(email_ids)
-                    print(f"{display_name:<40} {count:>10}")
-                    total_all_folders += count
-                else:
-                    print(f"{display_name:<40} {'Error':>10}")
-
-            except imaplib.IMAP4.error:
-                print(f"{display_name:<40} {'Error':>10}")
-
-        print("-" * 52)
-        print(f"{'TOTAL':<40} {total_all_folders:>10}")
-
-        # Logout
-        mail.logout()
-
-    except imaplib.IMAP4.error as e:
-        print(f"IMAP Error: {e}")
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-
-def _is_ignored_local_dir(dirname: str) -> bool:
-    return dirname.startswith(".") or dirname == "__pycache__"
-
-
-def list_local_folders(local_root: str) -> list[str]:
-    """List all folders under a local backup root in IMAP-style names."""
-    folders: set[str] = set()
-
-    for dirpath, dirnames, _filenames in os.walk(local_root):
-        dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)]
-
-        if os.path.abspath(dirpath) == os.path.abspath(local_root):
-            continue
-
-        rel = os.path.relpath(dirpath, local_root)
-        if rel == ".":
-            continue
-
-        parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)]
-        if not parts:
-            continue
-
-        folders.add("/".join(parts))
-
-    return sorted(folders)
-
-
-def get_local_email_count(local_root: str, folder_name: str):
-    """Return the count of .eml files in a local folder, or None if missing/unreadable."""
-    folder_path = os.path.join(local_root, *folder_name.split("/"))
-    if not os.path.isdir(folder_path):
-        return None
-
-    try:
-        count = 0
-        for filename in os.listdir(folder_path):
-            if not filename.endswith(".eml"):
-                continue
-            full_path = os.path.join(folder_path, filename)
-            if os.path.isfile(full_path):
-                count += 1
-        return count
-    except OSError:
-        return None
-
-
-def count_local_emails(local_path: str) -> None:
-    print(f"Scanning local backup: {local_path}")
-
-    folders = list_local_folders(local_path)
-    if not folders:
-        print("No folders found.")
-        return
-
-    total_all_folders = 0
-    print(f"{'Folder Name':<40} {'Count':>10}")
-    print("-" * 52)
-
-    for folder_name in folders:
-        count = get_local_email_count(local_path, folder_name)
-        if count is None:
-            print(f"{folder_name:<40} {'N/A':>10}")
-            continue
-
-        print(f"{folder_name:<40} {count:>10}")
-        total_all_folders += count
-
-    print("-" * 52)
-    print(f"{'TOTAL':<40} {total_all_folders:>10}")
-
-
-def main(argv: Optional[list[str]] = None) -> None:
-    # Phase 1: determine whether we're in local mode (--path)
-    default_path = os.getenv("BACKUP_LOCAL_PATH") or os.getenv("SRC_LOCAL_PATH")
-    pre_parser = argparse.ArgumentParser(add_help=False)
-    pre_parser.add_argument("--path", default=default_path)
-    pre_args, _ = pre_parser.parse_known_args(argv)
-    require_imap = not bool(pre_args.path)
-
-    # Phase 2: full parser with conditional requirements
-    parser = argparse.ArgumentParser(description="Count emails in IMAP account.")
-
-    parser.add_argument(
-        "--path",
-        default=default_path,
-        help="Local backup root to count (counts .eml files per folder). If set, IMAP args are ignored.",
-    )
-
-    # Try to unify var names for defaults. Priority: IMAP_* > SRC_IMAP_* > None
-    default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
-    default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
-    default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
-    default_client_id = os.getenv("OAUTH2_CLIENT_ID") or os.getenv("SRC_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--host",
-        default=default_host,
-        required=require_imap and not bool(default_host),
-        help="IMAP Server (or IMAP_HOST / SRC_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--user",
-        default=default_user,
-        required=require_imap and not bool(default_user),
-        help="Username (or IMAP_USERNAME / SRC_IMAP_USERNAME)",
-    )
-
-    auth_required = require_imap and not bool(default_pass or default_client_id)
-    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
-    auth_group.add_argument(
-        "--pass", dest="password", default=default_pass, help="Password (or IMAP_PASSWORD / SRC_IMAP_PASSWORD)"
-    )
-    auth_group.add_argument(
-        "--oauth2-client-id",
-        default=default_client_id,
-        dest="client_id",
-        help="OAuth2 Client ID (or OAUTH2_CLIENT_ID / SRC_OAUTH2_CLIENT_ID)",
-    )
-    auth_group.add_argument(
-        "--client-id",
-        default=default_client_id,
-        dest="client_id",
-        help=argparse.SUPPRESS,
-    )
-
-    parser.add_argument(
-        "--oauth2-client-secret",
-        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="client_secret",
-        help="OAuth2 Client Secret (if required) (or OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET)",
-    )
-    parser.add_argument(
-        "--client-secret",
-        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="client_secret",
-        help=argparse.SUPPRESS,
-    )
-
-    args = parser.parse_args(argv)
-
-    if args.path:
-        if not os.path.isdir(args.path):
-            print(f"Error: Local path does not exist or is not a directory: {args.path}")
-            sys.exit(1)
-
-        print("\n--- Configuration Summary ---")
-        print(f"Local Path      : {args.path}")
-        print("-----------------------------\n")
-        count_local_emails(args.path)
-        raise SystemExit(0)
-
-    IMAP_SERVER = args.host
-    USERNAME = args.user
-    PASSWORD = args.password
-
-    use_oauth2 = bool(args.client_id)
-
-    # Acquire OAuth2 token if configured
-    oauth2_token = None
-    oauth2_provider = None
-    if use_oauth2:
-        oauth2_provider = imap_oauth2.detect_oauth2_provider(IMAP_SERVER)
-        if not oauth2_provider:
-            print(f"Error: Could not detect OAuth2 provider from host '{IMAP_SERVER}'.")
-            sys.exit(1)
-        print(f"Acquiring OAuth2 token ({oauth2_provider})...")
-        oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
-            oauth2_provider, args.client_id, USERNAME, args.client_secret
-        )
-        if not oauth2_token:
-            print("Error: Failed to acquire OAuth2 token.")
-            sys.exit(1)
-        print("OAuth2 token acquired successfully.\n")
-
-    print("\n--- Configuration Summary ---")
-    print(f"Host            : {IMAP_SERVER}")
-    print(f"User            : {USERNAME}")
-    print(f"Auth Method     : {'OAuth2/' + oauth2_provider + ' (XOAUTH2)' if use_oauth2 else 'Basic (password)'}")
-    print("-----------------------------\n")
-
-    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
 
+from imap_count import main
 
 if __name__ == "__main__":
-    main()
+    print(
+        "WARNING: 'count_imap_emails.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_count.py'.",
+        file=sys.stderr,
+    )
+    print("Redirecting to 'imap_count.py'...", file=sys.stderr)
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_backup.py

@@ -0,0 +1,899 @@
+"""
+IMAP Email Backup Script
+
+Backs up emails from an IMAP account to a local directory.
+Stores each email as a separate .eml file (RFC 5322 format) which is compatible with
+most email clients (Thunderbird, Apple Mail, Outlook, etc.).
+
+Features:
+- Incremental Backup: Skips messages that have already been downloaded (checks existing UIDs locally).
+- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
+- Folder Replication: Recreates the IMAP folder structure locally.
+- Parallel Processing: Uses multithreading for fast downloads.
+- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
+
+Configuration (Environment Variables):
+    SRC_IMAP_HOST, SRC_IMAP_USERNAME: Source credentials.
+    SRC_IMAP_PASSWORD: Source password (or App Password).
+
+    OAuth2 (Optional - instead of password):
+    SRC_OAUTH2_CLIENT_ID: OAuth2 Client ID
+    SRC_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
+
+  BACKUP_LOCAL_PATH: Destination local directory.
+  MAX_WORKERS: Number of concurrent threads (default: 10).
+  BATCH_SIZE: Number of emails to process per batch (default: 10).
+  PRESERVE_LABELS: Set to "true" to create labels_manifest.json (Gmail). Default is "false".
+  PRESERVE_FLAGS: Set to "true" to preserve IMAP flags in manifest. Default is "false".
+  MANIFEST_ONLY: Set to "true" to only build manifest without downloading. Default is "false".
+  GMAIL_MODE: Set to "true" for Gmail backup mode. Default is "false".
+  DEST_DELETE: Set to "true" to delete local files not found on server (sync mode).
+              Default is "false".
+
+Usage:
+    python3 imap_backup.py \
+        --src-host "imap.example.com" \
+        --src-user "you@example.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup"
+
+Gmail Labels:
+    python3 imap_backup.py \
+        --src-host "imap.gmail.com" \
+        --src-user "you@gmail.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup" \
+        --preserve-labels \
+        "[Gmail]/All Mail"
+  This backs up all emails from [Gmail]/All Mail and creates a labels_manifest.json
+  file that maps each email's Message-ID to its Gmail labels for later restoration.
+"""
+
+import argparse
+import concurrent.futures
+import json
+import os
+import sys
+import threading
+
+from auth import imap_oauth2
+from core import imap_session
+from providers import provider_exchange, provider_gmail
+from utils import imap_common
+
+# Defaults
+MAX_WORKERS = 10
+BATCH_SIZE = 10
+MANIFEST_FILENAME = "labels_manifest.json"
+
+# Thread-local storage
+thread_local = threading.local()
+safe_print = imap_common.safe_print
+
+
+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):
+    src = imap_session.get_thread_connection(thread_local, "src", src_conf)
+    if not src:
+        safe_print("Error: Could not establish connection for batch.")
+        return
+
+    try:
+        src.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Error selecting folder {folder_name} in worker: {e}")
+        return
+
+    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, 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
+
+            success, src = process_single_uid(src, uid, folder_name, local_folder_path)
+            thread_local.src = src
+
+            if success:
+                break
+            if attempt < max_retries - 1:
+                src = None
+                thread_local.src = None
+
+
+def get_existing_uids(local_path):
+    """
+    Scans the local directory for files matching pattern matches {UID}_*.eml
+    Returns a set of UIDs (as strings).
+    """
+    existing = set()
+    if not os.path.exists(local_path):
+        return existing
+
+    try:
+        for filename in os.listdir(local_path):
+            if filename.endswith(".eml") and "_" in filename:
+                # Expecting UID_Subject.eml
+                parts = filename.split("_", 1)
+                if parts[0].isdigit():
+                    existing.add(parts[0])
+    except Exception:
+        pass
+    return existing
+
+
+# Standard IMAP flags that can be preserved during migration
+# \Recent is session-specific and cannot be set by clients
+# \Deleted should not be preserved as it marks messages for removal
+PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
+
+
+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,
+    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 = {}
+
+    # Ensure connection is healthy (refresh OAuth2 token if needed) before initial select
+    if src_conf:
+        imap_conn = imap_session.ensure_connection(imap_conn, src_conf)
+        if not imap_conn:
+            safe_print(f"Could not establish connection for folder {folder_name}")
+            return (message_info, None)
+
+    try:
+        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, imap_conn)
+
+    try:
+        resp, data = imap_conn.uid("search", None, "ALL")
+        if resp != "OK" or not data or not data[0]:
+            return (message_info, imap_conn)
+
+        uids = data[0].split()
+        if not uids:
+            return (message_info, imap_conn)
+
+        total_uids = len(uids)
+
+        # Fetch Message-IDs and FLAGS in batches - use larger batch for header-only fetches
+        batch_size = 200
+        for i in range(0, len(uids), batch_size):
+            batch = uids[i : i + batch_size]
+            uid_range = b",".join(batch)
+
+            # Report progress
+            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)])")
+                if resp != "OK":
+                    continue
+
+                # Parse response - items come in pairs for each message
+                for item in items:
+                    if isinstance(item, tuple) and len(item) >= 2:
+                        # First element contains UID and FLAGS info
+                        meta_str = (
+                            item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0])
+                        )
+
+                        # Extract all preservable flags from the metadata
+                        flags = []
+                        for flag in imap_common.PRESERVABLE_FLAGS:
+                            if flag in meta_str:
+                                flags.append(flag)
+
+                        # Second element contains the header
+                        msg_id = imap_common.extract_message_id(item[1])
+                        if msg_id:
+                            message_info[msg_id] = {"flags": flags}
+            except Exception as e:
+                safe_print(f"Error fetching batch in {folder_name}: {e}")
+                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.
+    This is a convenience wrapper around get_message_info_in_folder.
+    Optional progress_callback(current, total) for progress reporting.
+    """
+    info = get_message_info_in_folder(imap_conn, folder_name, progress_callback)
+    return set(info.keys())
+
+
+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
+    appear in each label, plus their flags from [Gmail]/All Mail.
+
+    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
+
+    manifest = {}
+    start_time = time.time()
+    total_emails_scanned = 0
+
+    safe_print("--- Building Gmail Labels Manifest ---")
+
+    # Get all selectable folders and filter to label folders
+    all_folders = imap_common.list_selectable_folders(imap_conn)
+    if not all_folders:
+        safe_print("Error: Could not list folders for label mapping.")
+        return manifest
+
+    # First, scan [Gmail]/All Mail to get the authoritative flags for all emails
+    safe_print("[0/N] Scanning [Gmail]/All Mail for flags (read/starred/etc)...")
+    all_mail_start = time.time()
+
+    def all_mail_progress_cb(current, total):
+        elapsed = time.time() - start_time
+        print(
+            f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+            end="",
+            flush=True,
+        )
+
+    all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
+        imap_conn, provider_gmail.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
+    )
+    print()  # New line after progress
+
+    # Initialize manifest with flags from All Mail
+    for msg_id, info in all_mail_info.items():
+        manifest[msg_id] = {"labels": [], "flags": info.get("flags", [])}
+
+    all_mail_elapsed = time.time() - all_mail_start
+    # Count flag statistics
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
+    safe_print(f"  -> {len(all_mail_info)} emails scanned ({all_mail_elapsed:.1f}s)")
+    safe_print(f"  -> Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}\n")
+
+    # Keep connection alive
+    try:
+        imap_conn.noop()
+    except Exception:
+        pass
+
+    # Parse folder names and filter to label folders
+    label_folders = [f for f in all_folders if provider_gmail.is_label_folder(f)]
+
+    total_folders = len(label_folders)
+    safe_print(f"Found {total_folders} label folders to scan.\n")
+
+    # Scan each label folder
+    for folder_idx, folder_name in enumerate(label_folders, 1):
+        folder_start = time.time()
+
+        # Progress callback for this folder
+        def progress_cb(current, total):
+            elapsed = time.time() - start_time
+            print(
+                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+                end="",
+                flush=True,
+            )
+
+        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
+        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
+        try:
+            imap_conn.noop()
+        except Exception:
+            pass
+
+        # Determine the label name to store
+        # For [Gmail]/Sent Mail -> "Sent Mail"
+        # For [Gmail]/Starred -> "Starred"
+        # For INBOX -> "INBOX"
+        # For user folders -> folder name as-is
+        if folder_name.startswith("[Gmail]/"):
+            label_name = folder_name[8:]  # Remove "[Gmail]/" prefix
+        else:
+            label_name = folder_name
+
+        for msg_id in message_ids:
+            if msg_id not in manifest:
+                # Email not in All Mail (rare, but handle it)
+                manifest[msg_id] = {"labels": [], "flags": []}
+            if label_name not in manifest[msg_id]["labels"]:
+                manifest[msg_id]["labels"].append(label_name)
+
+        folder_elapsed = time.time() - folder_start
+        total_emails_scanned += len(message_ids)
+        safe_print(f"  -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
+
+    # Summary
+    total_elapsed = time.time() - start_time
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
+    safe_print("\nManifest building complete:")
+    safe_print(f"  - Folders scanned: {total_folders + 1}")  # +1 for All Mail
+    safe_print(f"  - Total email-label mappings: {total_emails_scanned}")
+    safe_print(f"  - Unique emails: {len(manifest)}")
+    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}")
+    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
+
+    # Save manifest
+    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
+    try:
+        with open(manifest_path, "w", encoding="utf-8") as f:
+            json.dump(manifest, f, indent=2, ensure_ascii=False)
+        safe_print(f"\nLabels manifest saved to: {manifest_path}")
+    except Exception as e:
+        safe_print(f"Error saving manifest: {e}")
+
+    return manifest
+
+
+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
+
+    manifest = {}
+    start_time = time.time()
+
+    safe_print("--- Building Flags Manifest ---")
+
+    # Get folders to scan
+    if folders_to_scan:
+        all_folders = folders_to_scan
+    else:
+        try:
+            typ, folders = imap_conn.list()
+            if typ != "OK":
+                safe_print("Error: Could not list folders.")
+                return manifest
+            all_folders = [imap_common.normalize_folder_name(f) for f in folders]
+        except Exception as e:
+            safe_print(f"Error listing folders: {e}")
+            return manifest
+
+    total_folders = len(all_folders)
+    safe_print(f"Found {total_folders} folders to scan.\n")
+
+    # Scan each folder
+    for folder_idx, folder_name in enumerate(all_folders, 1):
+        folder_start = time.time()
+
+        def progress_cb(current, total):
+            elapsed = time.time() - start_time
+            print(
+                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+                end="",
+                flush=True,
+            )
+
+        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
+        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)
+        for msg_id, info in folder_info.items():
+            if msg_id not in manifest:
+                manifest[msg_id] = {"flags": info.get("flags", [])}
+
+        folder_elapsed = time.time() - folder_start
+        safe_print(f"  -> {len(folder_info)} emails ({folder_elapsed:.1f}s)")
+
+        # Keep connection alive
+        try:
+            imap_conn.noop()
+        except Exception:
+            pass
+
+    # Summary
+    total_elapsed = time.time() - start_time
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
+    safe_print("\nFlags manifest building complete:")
+    safe_print(f"  - Folders scanned: {total_folders}")
+    safe_print(f"  - Unique emails: {len(manifest)}")
+    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Flagged: {flagged_count}")
+    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
+
+    # Save manifest
+    manifest_path = os.path.join(local_path, "flags_manifest.json")
+    try:
+        with open(manifest_path, "w", encoding="utf-8") as f:
+            json.dump(manifest, f, indent=2, ensure_ascii=False)
+        safe_print(f"\nFlags manifest saved to: {manifest_path}")
+    except Exception as e:
+        safe_print(f"Error saving manifest: {e}")
+
+    return manifest
+
+
+def delete_orphan_local_files(local_folder_path, server_uids):
+    """
+    Delete local .eml files that no longer exist on the server.
+    Args:
+        local_folder_path: Path to local folder containing .eml files
+        server_uids: Set of UID strings currently on server
+    Returns:
+        Count of deleted files
+    """
+    deleted_count = 0
+    if not os.path.exists(local_folder_path):
+        return deleted_count
+
+    try:
+        for filename in os.listdir(local_folder_path):
+            if not filename.endswith(".eml") or "_" not in filename:
+                continue
+
+            # Extract UID from filename (format: {UID}_{Subject}.eml)
+            parts = filename.split("_", 1)
+            if not parts[0].isdigit():
+                continue
+
+            local_uid = parts[0]
+            if local_uid not in server_uids:
+                file_path = os.path.join(local_folder_path, filename)
+                try:
+                    os.remove(file_path)
+                    safe_print(f"  -> Deleted orphan: {filename}")
+                    deleted_count += 1
+                except Exception as e:
+                    safe_print(f"  -> Error deleting {filename}: {e}")
+    except Exception as e:
+        safe_print(f"Error scanning for orphan files: {e}")
+
+    return deleted_count
+
+
+def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=False):
+    safe_print(f"--- Processing Folder: {folder_name} ---")
+
+    # create local path
+    # Handle folder separators. IMAP output might be "Parent/Child"
+    # We rely on OS to handle "Parent/Child" as subdirectories using join
+    # But clean the segments
+    cleaned_name = folder_name.replace("/", os.sep)
+    local_folder_path = os.path.join(local_base_path, cleaned_name)
+
+    try:
+        os.makedirs(local_folder_path, exist_ok=True)
+    except Exception as e:
+        safe_print(f"Error creating directory {local_folder_path}: {e}")
+        return
+
+    # Select IMAP folder
+    try:
+        src_main.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Skipping {folder_name}: {e}")
+        return
+
+    # Search all
+    resp, data = src_main.uid("search", None, "ALL")
+    if resp != "OK":
+        return
+
+    uids = data[0].split()
+    total_on_server = len(uids)
+
+    # Build set of server UIDs for comparison
+    server_uid_set = set()
+    for u in uids:
+        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
+        server_uid_set.add(u_str)
+
+    if total_on_server == 0:
+        safe_print(f"Folder {folder_name} is empty.")
+        # If dest_delete enabled, delete all local files
+        if dest_delete:
+            deleted = delete_orphan_local_files(local_folder_path, set())
+            if deleted > 0:
+                safe_print(f"Deleted {deleted} orphan files from local backup.")
+        return
+
+    # Incremental Optimization
+    # Read local directory to find UIDs we already have
+    existing_uids = get_existing_uids(local_folder_path)
+
+    # Delete orphan local files if dest_delete is enabled
+    if dest_delete:
+        orphan_uids = existing_uids - server_uid_set
+        if orphan_uids:
+            safe_print(f"Found {len(orphan_uids)} local files not on server, deleting...")
+            deleted = delete_orphan_local_files(local_folder_path, server_uid_set)
+            if deleted > 0:
+                safe_print(f"Deleted {deleted} orphan files from local backup.")
+
+    # Filter UIDs
+    # decode uid first if bytes
+    uids_to_download = []
+    for u in uids:
+        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
+        if u_str not in existing_uids:
+            uids_to_download.append(u)
+
+    skipped = total_on_server - len(uids_to_download)
+    if skipped > 0:
+        safe_print(f"Skipping {skipped} emails (already exist locally).")
+
+    if not uids_to_download:
+        safe_print(f"Folder {folder_name} is up to date.")
+        return
+
+    safe_print(f"Downloading {len(uids_to_download)} new emails...")
+
+    # Create batches
+    uid_batches = [uids_to_download[i : i + BATCH_SIZE] for i in range(0, len(uids_to_download), BATCH_SIZE)]
+
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
+    try:
+        futures = []
+        for batch in uid_batches:
+            futures.append(executor.submit(process_batch, batch, folder_name, src_conf, local_folder_path))
+
+        for future in concurrent.futures.as_completed(futures):
+            future.result()
+
+    except KeyboardInterrupt:
+        executor.shutdown(wait=False, cancel_futures=True)
+        raise
+    finally:
+        executor.shutdown(wait=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+
+    # Source
+    default_src_host = os.getenv("SRC_IMAP_HOST")
+    default_src_user = os.getenv("SRC_IMAP_USERNAME")
+    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
+    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--src-host",
+        default=default_src_host,
+        required=not bool(default_src_host),
+        help="Source IMAP Server (or SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--src-user",
+        default=default_src_user,
+        required=not bool(default_src_user),
+        help="Source Username (or SRC_IMAP_USERNAME)",
+    )
+
+    # Authentication: require either password OR OAuth2 client-id (unless provided via env vars)
+    auth_required = not bool(default_src_pass or default_src_client_id)
+    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
+    auth_group.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
+    # OAuth2
+    auth_group.add_argument(
+        "--src-oauth2-client-id",
+        default=default_src_client_id,
+        dest="src_client_id",
+        help="OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
+    )
+    parser.add_argument(
+        "--src-oauth2-client-secret",
+        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="src_client_secret",
+        help="OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
+    )
+
+    # Destination (Local Path)
+    env_path = os.getenv("BACKUP_LOCAL_PATH")
+    parser.add_argument(
+        "--dest-path",
+        default=env_path,
+        required=not bool(env_path),
+        help="Local destination path (or BACKUP_LOCAL_PATH)",
+    )
+
+    # Config
+    parser.add_argument("--workers", type=int, default=int(os.getenv("MAX_WORKERS", 10)), help="Thread count")
+    parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
+
+    # Gmail Labels
+    env_preserve_labels = os.getenv("PRESERVE_LABELS", "false").lower() == "true"
+    parser.add_argument(
+        "--preserve-labels",
+        action="store_true",
+        default=env_preserve_labels,
+        help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
+    )
+    env_preserve_flags = os.getenv("PRESERVE_FLAGS", "false").lower() == "true"
+    parser.add_argument(
+        "--preserve-flags",
+        action="store_true",
+        default=env_preserve_flags,
+        help="Preserve IMAP flags (read/unread, starred, answered, draft) in manifest for restoration",
+    )
+    env_manifest_only = os.getenv("MANIFEST_ONLY", "false").lower() == "true"
+    parser.add_argument(
+        "--manifest-only",
+        action="store_true",
+        default=env_manifest_only,
+        help="Gmail only: Build the labels manifest and exit without downloading emails",
+    )
+    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
+    parser.add_argument(
+        "--gmail-mode",
+        action="store_true",
+        default=env_gmail_mode,
+        help="Gmail backup mode: Build labels manifest and backup [Gmail]/All Mail only (recommended)",
+    )
+
+    # Sync mode: delete local files not on server
+    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
+    parser.add_argument(
+        "--dest-delete",
+        action="store_true",
+        default=env_dest_delete,
+        help="Delete local .eml files that no longer exist on the IMAP server (sync mode)",
+    )
+
+    parser.add_argument("folder", nargs="?", help="Specific folder to backup")
+
+    args = parser.parse_args()
+
+    global MAX_WORKERS, BATCH_SIZE
+    MAX_WORKERS = args.workers
+    BATCH_SIZE = args.batch
+
+    # Build connection config (acquires OAuth2 token if configured)
+    src_conf = imap_session.build_imap_conf(
+        args.src_host, args.src_user, args.src_pass, args.src_client_id, args.src_client_secret
+    )
+
+    # Expand path (~/...)
+    local_path = os.path.expanduser(args.dest_path)
+    if not os.path.exists(local_path):
+        try:
+            os.makedirs(local_path)
+            print(f"Created backup directory: {local_path}")
+        except Exception as e:
+            print(f"Error creating backup directory: {e}")
+            sys.exit(1)
+
+    print("\n--- Configuration Summary ---")
+    print(f"Source Host     : {args.src_host}")
+    print(f"Source User     : {args.src_user}")
+    print(f"Auth Method     : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
+    print(f"Destination Path: {local_path}")
+    if args.gmail_mode:
+        print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
+    elif args.manifest_only:
+        print("Mode            : Manifest Only (no email download)")
+    elif args.folder:
+        print(f"Target Folder   : {args.folder}")
+    if args.preserve_labels or args.manifest_only or args.gmail_mode:
+        print("Preserve Labels : Yes (Gmail)")
+    if args.preserve_flags or args.gmail_mode:
+        print("Preserve Flags  : Yes (read/starred/answered/draft)")
+    if args.dest_delete:
+        print("Dest Delete     : Yes (remove local orphans)")
+    print("-----------------------------\n")
+
+    try:
+        src = imap_common.get_imap_connection_from_conf(src_conf)
+        if not src:
+            sys.exit(1)
+
+        # Build labels manifest BEFORE backing up emails (Gmail mode)
+        # This way we capture the label state at backup time
+        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, src_conf)
+            print("")  # Blank line after manifest building
+        # Build flags-only manifest for non-Gmail servers
+        elif args.preserve_flags:
+            print("Building flags manifest...")
+            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, 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}")
+            print("\nTo download emails, run again without --manifest-only:")
+            print(f'  python3 imap_backup.py --dest-path "{local_path}" "[Gmail]/All Mail"')
+            sys.exit(0)
+
+        # Gmail mode: backup only [Gmail]/All Mail
+        if args.gmail_mode:
+            backup_folder(src, provider_gmail.GMAIL_ALL_MAIL, local_path, src_conf, args.dest_delete)
+        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:
+                if provider_exchange.is_special_folder(name):
+                    print(f"Skipping Exchange system folder: {name}")
+                    continue
+                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:
+            manifest_path = os.path.join(local_path, "labels_manifest.json")
+            print(f"\nGmail labels manifest saved to: {manifest_path}")
+            print("Use this file when restoring to reapply labels and flags to emails.")
+        elif args.preserve_flags:
+            manifest_path = os.path.join(local_path, "flags_manifest.json")
+            print(f"\nFlags manifest saved to: {manifest_path}")
+            print("Use this file when restoring to reapply read/starred status to emails.")
+
+    except KeyboardInterrupt:
+        raise
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_common.py

@@ -1,438 +0,0 @@
-"""
-IMAP Common Utilities
-
-Shared functionality for IMAP migration, counting, and comparison scripts.
-"""
-
-from __future__ import annotations
-
-import imaplib
-import os
-import re
-import sys
-from email import policy
-from email.header import decode_header
-from email.parser import BytesParser
-
-# Standard IMAP flags
-FLAG_SEEN = "\\Seen"
-FLAG_ANSWERED = "\\Answered"
-FLAG_FLAGGED = "\\Flagged"
-FLAG_DRAFT = "\\Draft"
-FLAG_DELETED = "\\Deleted"
-FLAG_DELETED_LITERAL = "(\\Deleted)"
-
-# Standard IMAP flags that can be preserved during migration
-# \Recent is session-specific and cannot be set by clients
-# \Deleted should not be preserved as it marks messages for removal
-PRESERVABLE_FLAGS = {FLAG_SEEN, FLAG_ANSWERED, FLAG_FLAGGED, FLAG_DRAFT}
-
-# Gmail constants
-GMAIL_ALL_MAIL = "[Gmail]/All Mail"
-GMAIL_TRASH = "[Gmail]/Trash"
-GMAIL_SPAM = "[Gmail]/Spam"
-GMAIL_DRAFTS = "[Gmail]/Drafts"
-GMAIL_BIN = "[Gmail]/Bin"
-GMAIL_IMPORTANT = "[Gmail]/Important"
-GMAIL_SENT = "[Gmail]/Sent Mail"
-GMAIL_STARRED = "[Gmail]/Starred"
-
-GMAIL_SYSTEM_FOLDERS = {
-    GMAIL_ALL_MAIL,
-    GMAIL_SPAM,
-    GMAIL_TRASH,
-    GMAIL_DRAFTS,
-    GMAIL_BIN,
-    GMAIL_IMPORTANT,
-}
-
-# IMAP Folder Constants
-FOLDER_INBOX = "INBOX"
-
-# Gmail-mode restore/migrate fallback folder for messages with no usable labels.
-# Keeping this as a normal folder (not [Gmail]/Drafts) avoids populating Drafts with non-drafts.
-FOLDER_RESTORED_UNLABELED = "Restored/Unlabeled"
-
-# IMAP Commands
-CMD_STORE = "store"
-CMD_SEARCH = "search"
-CMD_FETCH = "fetch"
-OP_ADD_FLAGS = "+FLAGS"
-
-
-def ensure_folder_exists(imap_conn, folder_name: str) -> None:
-    """Best-effort create of a folder if it doesn't already exist.
-
-    IMAP servers typically return an error if the mailbox exists; this helper
-    intentionally ignores those errors.
-    """
-    try:
-        if folder_name and folder_name.upper() != FOLDER_INBOX:
-            imap_conn.create(f'"{folder_name}"')
-    except Exception:
-        # Folder may already exist, or server may restrict creation.
-        pass
-
-
-def append_email(
-    imap_conn,
-    folder_name: str,
-    raw_content: bytes,
-    date_str: str,
-    flags: str | None = None,
-    *,
-    ensure_folder: bool = True,
-) -> bool:
-    """Append an email message to a folder.
-
-    This is intentionally a thin wrapper around IMAP APPEND; callers can
-    perform duplicate checks or folder selection separately.
-
-    Args:
-        flags: Optional IMAP flags. If provided, they are normalized to a
-            parenthesized list before being passed to `imaplib.IMAP4.append`.
-        ensure_folder: If True, attempts to create the folder first (best-effort).
-    """
-    try:
-        if ensure_folder:
-            ensure_folder_exists(imap_conn, folder_name)
-
-        normalized_flags = None
-        if flags:
-            stripped = str(flags).strip()
-            if stripped:
-                if stripped.startswith("(") and stripped.endswith(")"):
-                    normalized_flags = stripped
-                else:
-                    normalized_flags = f"({stripped})"
-
-        resp, _ = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
-        return resp == "OK"
-    except Exception:
-        return False
-
-
-def verify_env_vars(vars_list):
-    """
-    Checks if all environment variables in the list are set.
-    Returns True if all are present, False otherwise.
-    Prints missing variables to stderr.
-    """
-    missing = [v for v in vars_list if not os.getenv(v)]
-    if missing:
-        print(f"Error: Missing environment variables: {', '.join(missing)}", file=sys.stderr)
-        return False
-    return True
-
-
-def get_imap_connection_from_conf(conf):
-    """
-    Establishes an IMAP connection using a conf dict.
-
-    conf dict structure:
-        {
-            "host": str,
-            "user": str,
-            "password": str or None,
-            "oauth2_token": str or None,
-            "oauth2": dict or None  # Contains provider, client_id, email, client_secret
-        }
-    """
-    return get_imap_connection(conf["host"], conf["user"], conf.get("password"), conf.get("oauth2_token"))
-
-
-def get_imap_connection(host, user, password=None, oauth2_token=None):
-    """
-    Establishes an SSL connection to the IMAP server and logs in.
-    Supports both basic auth (password) and OAuth 2.0 (XOAUTH2).
-    Returns the connection object or None if failed.
-    """
-    if not host or not user:
-        print(f"Error: Invalid credentials for {host}")
-        return None
-
-    if not password and not oauth2_token:
-        print(f"Error: Either password or oauth2_token is required for {host}")
-        return None
-
-    try:
-        conn = imaplib.IMAP4_SSL(host)
-        if oauth2_token:
-            auth_string = f"user={user}\x01auth=Bearer {oauth2_token}\x01\x01"
-            conn.authenticate("XOAUTH2", lambda _: auth_string.encode())
-        else:
-            conn.login(user, password)
-        return conn
-    except Exception as e:
-        print(f"Connection error to {host}: {e}")
-        return None
-
-
-def normalize_folder_name(folder_info_str):
-    """
-    Parses the IMAP list response to extract the clean folder name.
-    Handles quoted names and flags.
-    """
-    if isinstance(folder_info_str, bytes):
-        folder_info_str = folder_info_str.decode("utf-8", errors="ignore")
-
-    # Regex to extract folder name: (flags) "delimiter" name
-    # Matches: (\HasNoChildren) "/" "INBOX"  OR  (\HasNoChildren) "/" Drafts
-    list_pattern = re.compile(r'\((?P<flags>.*?)\) "(?P<delimiter>.*)" "?(?P<name>.*)"?')
-    match = list_pattern.search(folder_info_str)
-    if match:
-        name = match.group("name")
-        # If the regex grabbed a trailing quote, strip it (though the regex tries to handle it)
-        return name.rstrip('"').strip()
-
-    # Fallback: take the last part
-    return folder_info_str.split()[-1].strip('"')
-
-
-def list_selectable_folders(imap_conn):
-    """
-    Lists all selectable folders (excluding \\Noselect) on the IMAP connection.
-    Returns a list of normalized folder name strings, or an empty list on failure.
-    """
-    try:
-        status, folders = imap_conn.list()
-        if status != "OK" or not folders:
-            return []
-    except Exception:
-        return []
-
-    result = []
-    for f in folders:
-        f_str = f.decode("utf-8", errors="ignore") if isinstance(f, bytes) else str(f)
-        if "\\Noselect" in f_str:
-            continue
-        result.append(normalize_folder_name(f))
-    return result
-
-
-def decode_mime_header(header_value):
-    """
-    Decodes MIME encoded headers (Subject, etc.) to a unicode string.
-    """
-    if not header_value:
-        return "(No Subject)"
-    try:
-        decoded_list = decode_header(header_value)
-        text_parts = []
-        for data, encoding in decoded_list:
-            if isinstance(data, bytes):
-                charset = encoding or "utf-8"
-                try:
-                    text_parts.append(data.decode(charset, errors="ignore"))
-                except LookupError:
-                    text_parts.append(data.decode("utf-8", errors="ignore"))
-            else:
-                text_parts.append(str(data))
-        return "".join(text_parts)
-    except Exception:
-        return str(header_value)
-
-
-def decode_message_id(msg_id):
-    """
-    Decodes a Message-ID header value by unfolding continuation lines.
-    Returns the stripped Message-ID string or None if empty.
-    """
-    if not msg_id:
-        return None
-    # Unfold any header continuation lines (CRLF/LF + whitespace)
-    return re.sub(r"\r?\n[ \t]+", " ", str(msg_id)).strip() or None
-
-
-def extract_message_id(header_data):
-    """
-    Extracts the Message-ID from header bytes or string using BytesParser.
-    Returns the stripped Message-ID string or None if not found.
-    """
-    if not header_data:
-        return None
-
-    try:
-        # Use compat32 to preserve raw header with continuation lines
-        # (policy.default's _MessageIDHeader parser truncates at newlines)
-        parser = BytesParser(policy=policy.compat32)
-        if isinstance(header_data, str):
-            header_data = header_data.encode("utf-8", errors="ignore")
-
-        msg = parser.parsebytes(header_data, headersonly=True)
-        return decode_message_id(msg.get("Message-ID"))
-    except Exception:
-        pass
-    return None
-
-
-def parse_message_id_from_bytes(raw_message):
-    """Parse Message-ID from RFC822 bytes.
-
-    Parses the full message bytes to extract the Message-ID header.
-    """
-    if not raw_message:
-        return None
-    try:
-        parser = BytesParser()
-        msg = parser.parsebytes(raw_message)
-        return decode_message_id(msg.get("Message-ID"))
-    except Exception:
-        return None
-
-
-def parse_message_id_and_subject_from_bytes(raw_message):
-    """Parse Message-ID and decoded Subject from RFC822 bytes.
-
-    Uses header-only parsing to avoid walking large message bodies.
-    Returns a tuple: (message_id, subject).
-    """
-    if not raw_message:
-        return None, "(No Subject)"
-
-    try:
-        # Use compat32 to preserve raw headers with continuation lines
-        parser = BytesParser(policy=policy.compat32)
-        email_obj = parser.parsebytes(raw_message, headersonly=True)
-        msg_id = decode_message_id(email_obj.get("Message-ID"))
-        raw_subject = email_obj.get("Subject")
-        subject = decode_mime_header(raw_subject) if raw_subject else "(No Subject)"
-        return msg_id, subject
-    except Exception:
-        return None, "(No 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.
-    Returns True if found, False otherwise.
-    """
-    if not msg_id:
-        return False
-
-    clean_id = msg_id.replace('"', '\\"')
-    try:
-        typ, data = dest_conn.search(None, f'(HEADER Message-ID "{clean_id}")')
-        if typ != "OK":
-            return False
-
-        dest_ids = data[0].split()
-        return len(dest_ids) > 0
-    except Exception:
-        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.
-    Removes/replaces characters that are illegal in file systems.
-    Truncates to 250 chars.
-    """
-    if not filename:
-        return "untitled"
-    # Replace invalid characters with underscore
-    # Invalid: < > : " / \ | ? * and control chars
-    s = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", filename)
-    # Strip leading/trailing whitespaces/dots
-    s = s.strip().strip(".")
-    # Ensure not empty and not too long
-    return s[:250] if s else "untitled"
-
-
-def detect_trash_folder(imap_conn):
-    """
-    Attempts to identify the Trash folder in the account.
-    Returns the folder name (str) or None if not found.
-    Checks for common names and SPECIAL-USE attributes.
-    """
-    try:
-        status, folders = imap_conn.list()
-        if status != "OK":
-            return None
-    except Exception:
-        return None
-
-    trash_candidates = ["[Gmail]/Trash", "Trash", "Deleted Items", "Bin", "[Gmail]/Bin"]
-    detected_by_flag = None
-    all_folder_names = []
-
-    for f in folders:
-        if isinstance(f, bytes):
-            f_str = f.decode("utf-8", errors="ignore")
-        else:
-            f_str = str(f)
-
-        name = normalize_folder_name(f_str)
-        all_folder_names.append(name)
-
-        # Check for SPECIAL-USE flag \Trash
-        # The flag is usually inside parentheses like (\HasNoChildren \Trash)
-        if "\\Trash" in f_str or "\\Bin" in f_str:
-            detected_by_flag = name
-
-    if detected_by_flag:
-        return detected_by_flag
-
-    # Check candidates
-    for candidate in trash_candidates:
-        if candidate in all_folder_names:
-            return candidate
-
-    return None

src/imap_compare.py

@@ -0,0 +1,337 @@
+"""
+IMAP Folder Comparison Script
+
+This script compares email counts between a source and a destination.
+Each side can be either an IMAP account or a local backup folder.
+It iterates through all folders found in the source account and checks the corresponding
+folder in the destination account.
+
+Configuration (Environment Variables):
+  Source Account:
+    SRC_IMAP_HOST       : Source IMAP Host
+    SRC_IMAP_USERNAME   : Source Username/Email
+    SRC_IMAP_PASSWORD   : Source Password
+
+    OAuth2 (Optional - instead of password):
+    SRC_OAUTH2_CLIENT_ID     : OAuth2 Client ID
+    SRC_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
+
+  Destination Account:
+    DEST_IMAP_HOST      : Destination IMAP Host
+    DEST_IMAP_USERNAME  : Destination Username/Email
+    DEST_IMAP_PASSWORD  : Destination Password
+
+    OAuth2 (Optional - instead of password):
+    DEST_OAUTH2_CLIENT_ID     : OAuth2 Client ID
+    DEST_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
+
+Also supports local folders as source and/or destination:
+    SRC_LOCAL_PATH      : Source local folder (backup root)
+    DEST_LOCAL_PATH     : Destination local folder (backup root)
+
+Usage:
+  python3 imap_compare.py
+
+Examples:
+        # IMAP -> IMAP
+        python3 imap_compare.py \
+            --src-host "imap.source.com" \
+            --src-user "source@example.com" \
+            --src-pass "source-app-password" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # Local -> IMAP
+        python3 imap_compare.py \
+            --src-path "./my_backup" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # IMAP -> Local
+        python3 imap_compare.py \
+            --src-host "imap.source.com" \
+            --src-user "source@example.com" \
+            --src-pass "source-app-password" \
+            --dest-path "./my_backup"
+"""
+
+import argparse
+import os
+import sys
+
+from auth import imap_oauth2
+from utils import imap_common
+
+
+def get_email_count(conn, folder_name):
+    """Return the IMAP message count for a folder, or None on error."""
+    try:
+        # Select folder in read-only mode
+        # Quote folder name handles spaces
+        typ, data = conn.select(f'"{folder_name}"', readonly=True)
+        if typ != "OK":
+            return None
+
+        # SELECT command returns the number of messages in data[0]
+        # data[0] is bytes, e.g. b'123'
+        if data and data[0]:
+            return int(data[0])
+        return 0
+
+    except Exception:
+        # print(f"Error checking {folder_name}: {e}")
+        return None
+
+
+def main():
+    default_src_path = os.getenv("SRC_LOCAL_PATH")
+    default_dest_path = os.getenv("DEST_LOCAL_PATH")
+
+    # Phase 1: determine whether each side is local
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--src-path", default=default_src_path)
+    pre_parser.add_argument("--dest-path", default=default_dest_path)
+    pre_args, _ = pre_parser.parse_known_args()
+
+    src_requires_imap = not bool(pre_args.src_path)
+    dest_requires_imap = not bool(pre_args.dest_path)
+
+    # Phase 2: full parser with conditional requirements
+    parser = argparse.ArgumentParser(description="Compare email counts between two IMAP accounts.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+    parser.add_argument(
+        "--src-path",
+        default=default_src_path,
+        help="Source local folder (backup root). If set, IMAP source args are ignored.",
+    )
+    parser.add_argument(
+        "--dest-path",
+        default=default_dest_path,
+        help="Destination local folder (backup root). If set, IMAP destination args are ignored.",
+    )
+
+    # Source args
+    default_src_host = os.getenv("SRC_IMAP_HOST")
+    default_src_user = os.getenv("SRC_IMAP_USERNAME")
+    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
+    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--src-host",
+        default=default_src_host,
+        required=src_requires_imap and not bool(default_src_host),
+        help="Source IMAP Server (or SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--src-user",
+        default=default_src_user,
+        required=src_requires_imap and not bool(default_src_user),
+        help="Source Username (or SRC_IMAP_USERNAME)",
+    )
+    src_auth_required = src_requires_imap and not bool(default_src_pass or default_src_client_id)
+    src_auth = parser.add_mutually_exclusive_group(required=src_auth_required)
+    src_auth.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
+    src_auth.add_argument(
+        "--src-oauth2-client-id",
+        default=default_src_client_id,
+        dest="src_client_id",
+        help="Source OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
+    )
+    parser.add_argument(
+        "--src-oauth2-client-secret",
+        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="src_client_secret",
+        help="Source OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
+    )
+
+    # Dest args
+    default_dest_host = os.getenv("DEST_IMAP_HOST")
+    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
+    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
+    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--dest-host",
+        default=default_dest_host,
+        required=dest_requires_imap and not bool(default_dest_host),
+        help="Destination IMAP Server (or DEST_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--dest-user",
+        default=default_dest_user,
+        required=dest_requires_imap and not bool(default_dest_user),
+        help="Destination Username (or DEST_IMAP_USERNAME)",
+    )
+    dest_auth_required = dest_requires_imap and not bool(default_dest_pass or default_dest_client_id)
+    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
+    dest_auth.add_argument(
+        "--dest-pass",
+        default=default_dest_pass,
+        help="Destination Password (or DEST_IMAP_PASSWORD)",
+    )
+    dest_auth.add_argument(
+        "--dest-oauth2-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
+    )
+    dest_auth.add_argument(
+        "--dest-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--dest-oauth2-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--dest-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help=argparse.SUPPRESS,
+    )
+
+    args = parser.parse_args()
+
+    src_is_local = bool(args.src_path)
+    dest_is_local = bool(args.dest_path)
+
+    SRC_HOST = args.src_host
+    SRC_USER = args.src_user
+    DEST_HOST = args.dest_host
+    DEST_USER = args.dest_user
+
+    # Acquire OAuth2 tokens if configured
+    src_oauth2_token = None
+    src_oauth2_provider = None
+    if not src_is_local and args.src_client_id:
+        src_oauth2_token, src_oauth2_provider = imap_oauth2.acquire_token(
+            SRC_HOST, args.src_client_id, SRC_USER, args.src_client_secret, "source"
+        )
+
+    dest_oauth2_token = None
+    dest_oauth2_provider = None
+    if not dest_is_local and args.dest_client_id:
+        dest_oauth2_token, dest_oauth2_provider = imap_oauth2.acquire_token(
+            DEST_HOST, args.dest_client_id, DEST_USER, args.dest_client_secret, "destination"
+        )
+
+    print("\n--- Configuration Summary ---")
+    if src_is_local:
+        print(f"Source (Local)  : {args.src_path}")
+    else:
+        print(f"Source Host     : {args.src_host}")
+        print(f"Source User     : {args.src_user}")
+        print(f"Source Auth     : {imap_oauth2.auth_description(src_oauth2_provider)}")
+
+    if dest_is_local:
+        print(f"Destination (Local): {args.dest_path}")
+    else:
+        print(f"Destination Host: {args.dest_host}")
+        print(f"Destination User: {args.dest_user}")
+        print(f"Destination Auth: {imap_oauth2.auth_description(dest_oauth2_provider)}")
+    print("-----------------------------\n")
+
+    src = None
+    dest = None
+
+    try:
+        if not src_is_local:
+            # Connect to Source
+            print("Connecting to Source...")
+            src = imap_common.get_imap_connection(args.src_host, args.src_user, args.src_pass, src_oauth2_token)
+            if not src:
+                return
+
+        if not dest_is_local:
+            # Connect to Dest
+            print("Connecting to Destination...")
+            dest = imap_common.get_imap_connection(args.dest_host, args.dest_user, args.dest_pass, dest_oauth2_token)
+            if not dest:
+                return
+
+        # List Source Folders
+        print("Listing folders in Source...")
+        if src_is_local:
+            folders = imap_common.list_local_folders(args.src_path)
+        else:
+            folders = imap_common.list_selectable_folders(src)
+
+        if not folders:
+            print("Failed to list source folders.")
+            return
+
+        # Prepare Table Header
+        header = f"{'Folder Name':<40} | {'Source':>10} | {'Dest':>10} | {'Diff':>10}"
+        print("-" * len(header))
+        print(header)
+        print("-" * len(header))
+
+        total_src = 0
+        total_dest = 0
+
+        # Iterate through Source folders
+        for folder_name in folders:
+            # Get Counts
+            if src_is_local:
+                src_count = imap_common.get_local_email_count(args.src_path, folder_name)
+            else:
+                src_count = get_email_count(src, folder_name)
+
+            if dest_is_local:
+                dest_count = imap_common.get_local_email_count(args.dest_path, folder_name)
+            else:
+                dest_count = get_email_count(dest, folder_name)
+
+            # Format for display
+            src_str = str(src_count) if src_count is not None else "Err"
+            dest_str = str(dest_count) if dest_count is not None else "N/A"  # N/A usually means folder doesn't exist
+
+            diff_str = ""
+            if src_count is not None and dest_count is not None:
+                diff = src_count - dest_count
+                diff_str = str(diff)
+                total_src += src_count
+                total_dest += dest_count
+            elif src_count is not None:
+                total_src += src_count
+
+            print(f"{folder_name:<40} | {src_str:>10} | {dest_str:>10} | {diff_str:>10}")
+
+        print("-" * len(header))
+        print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
+
+    except KeyboardInterrupt:
+        # Re-raise to be handled by the outer block, but ensure finally runs
+        raise
+
+    finally:
+        # Check source connection state and logout if possible
+        if src:
+            try:
+                src.logout()
+            except BaseException:
+                pass
+
+        # Check dest connection state and logout if possible
+        if dest:
+            try:
+                dest.logout()
+            except BaseException:
+                pass
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_count.py

@@ -0,0 +1,247 @@
+"""IMAP Email Counting Script.
+
+Counts emails per folder from either:
+- An IMAP account, or
+- A local backup folder created by ``imap_backup.py`` (counts ``.eml`` files).
+
+Configuration (Environment Variables):
+    IMAP_HOST        : IMAP Host (e.g., imap.gmail.com)
+    IMAP_USERNAME    : Username/Email
+    IMAP_PASSWORD    : Password (or App Password)
+
+    OAuth2 (Optional - instead of password):
+    OAUTH2_CLIENT_ID        : OAuth2 Client ID
+    OAUTH2_CLIENT_SECRET    : OAuth2 Client Secret (required for Google)
+    SRC_OAUTH2_CLIENT_ID    : Alternate OAuth2 client ID env var
+    SRC_OAUTH2_CLIENT_SECRET: Alternate OAuth2 client secret env var
+
+Local backup counting:
+    BACKUP_LOCAL_PATH : Local backup root (preferred)
+    SRC_LOCAL_PATH    : Alternate local backup root
+
+Examples:
+    # Count an IMAP account
+    export IMAP_HOST="imap.gmail.com"
+    export IMAP_USERNAME="user@gmail.com"
+    export IMAP_PASSWORD="secretpassword"
+    python3 imap_count.py
+
+    # Count an IMAP account using OAuth2
+    export IMAP_HOST="imap.gmail.com"
+    export IMAP_USERNAME="user@gmail.com"
+    export OAUTH2_CLIENT_ID="your-client-id"
+    export OAUTH2_CLIENT_SECRET="your-client-secret"  # Required for Google
+    python3 imap_count.py
+
+    # Count a local backup
+    python3 imap_count.py --path "./my_backup"
+
+    # Or set a default local backup path via env var
+    export BACKUP_LOCAL_PATH="./my_backup"
+    python3 imap_count.py
+"""
+
+import argparse
+import imaplib
+import os
+import sys
+from typing import Optional
+
+from auth import imap_oauth2
+from utils import imap_common
+
+
+def count_emails(imap_server, username, password=None, oauth2_token=None):
+    try:
+        # Connect to the IMAP server (using SSL)
+        print(f"Connecting to {imap_server}...")
+        mail = imap_common.get_imap_connection(imap_server, username, password, oauth2_token)
+        if not mail:
+            return
+
+        # List all mailboxes
+        print("Listing mailboxes...")
+        folders = imap_common.list_selectable_folders(mail)
+
+        if not folders:
+            print("Failed to list mailboxes.")
+            return
+
+        total_all_folders = 0
+        print(f"{'Folder Name':<40} {'Count':>10}")
+        print("-" * 52)
+
+        for folder_name in folders:
+            display_name = folder_name
+
+            try:
+                # Select the mailbox (read-only is sufficient for counting)
+                # folder_name extracted from list usually handles quotes correctly for select
+                rv, _ = mail.select(f'"{folder_name}"', readonly=True)
+                if rv != "OK":
+                    print(f"{display_name:<40} {'Skipped':>10}")
+                    continue
+
+                # Search for all emails
+                status, data = mail.search(None, "ALL")
+
+                if status == "OK":
+                    # data[0] is space separated IDs
+                    email_ids = data[0].split()
+                    count = len(email_ids)
+                    print(f"{display_name:<40} {count:>10}")
+                    total_all_folders += count
+                else:
+                    print(f"{display_name:<40} {'Error':>10}")
+
+            except imaplib.IMAP4.error:
+                print(f"{display_name:<40} {'Error':>10}")
+
+        print("-" * 52)
+        print(f"{'TOTAL':<40} {total_all_folders:>10}")
+
+        # Logout
+        mail.logout()
+
+    except imaplib.IMAP4.error as e:
+        print(f"IMAP Error: {e}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+def count_local_emails(local_path: str) -> None:
+    print(f"Scanning local backup: {local_path}")
+
+    folders = imap_common.list_local_folders(local_path)
+    if not folders:
+        print("No folders found.")
+        return
+
+    total_all_folders = 0
+    print(f"{'Folder Name':<40} {'Count':>10}")
+    print("-" * 52)
+
+    for folder_name in folders:
+        count = imap_common.get_local_email_count(local_path, folder_name)
+        if count is None:
+            print(f"{folder_name:<40} {'N/A':>10}")
+            continue
+
+        print(f"{folder_name:<40} {count:>10}")
+        total_all_folders += count
+
+    print("-" * 52)
+    print(f"{'TOTAL':<40} {total_all_folders:>10}")
+
+
+def main(argv: Optional[list[str]] = None) -> None:
+    # Phase 1: determine whether we're in local mode (--path)
+    default_path = os.getenv("BACKUP_LOCAL_PATH") or os.getenv("SRC_LOCAL_PATH")
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--path", default=default_path)
+    pre_args, _ = pre_parser.parse_known_args(argv)
+    require_imap = not bool(pre_args.path)
+
+    # Phase 2: full parser with conditional requirements
+    parser = argparse.ArgumentParser(description="Count emails in IMAP account.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+
+    parser.add_argument(
+        "--path",
+        default=default_path,
+        help="Local backup root to count (counts .eml files per folder). If set, IMAP args are ignored.",
+    )
+
+    # Try to unify var names for defaults. Priority: IMAP_* > SRC_IMAP_* > None
+    default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
+    default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
+    default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
+    default_client_id = os.getenv("OAUTH2_CLIENT_ID") or os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--host",
+        default=default_host,
+        required=require_imap and not bool(default_host),
+        help="IMAP Server (or IMAP_HOST / SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--user",
+        default=default_user,
+        required=require_imap and not bool(default_user),
+        help="Username (or IMAP_USERNAME / SRC_IMAP_USERNAME)",
+    )
+
+    auth_required = require_imap and not bool(default_pass or default_client_id)
+    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
+    auth_group.add_argument(
+        "--pass", dest="password", default=default_pass, help="Password (or IMAP_PASSWORD / SRC_IMAP_PASSWORD)"
+    )
+    auth_group.add_argument(
+        "--oauth2-client-id",
+        default=default_client_id,
+        dest="client_id",
+        help="OAuth2 Client ID (or OAUTH2_CLIENT_ID / SRC_OAUTH2_CLIENT_ID)",
+    )
+    auth_group.add_argument(
+        "--client-id",
+        default=default_client_id,
+        dest="client_id",
+        help=argparse.SUPPRESS,
+    )
+
+    parser.add_argument(
+        "--oauth2-client-secret",
+        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="client_secret",
+        help="OAuth2 Client Secret (if required) (or OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--client-secret",
+        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="client_secret",
+        help=argparse.SUPPRESS,
+    )
+
+    args = parser.parse_args(argv)
+
+    if args.path:
+        if not os.path.isdir(args.path):
+            print(f"Error: Local path does not exist or is not a directory: {args.path}")
+            sys.exit(1)
+
+        print("\n--- Configuration Summary ---")
+        print(f"Local Path      : {args.path}")
+        print("-----------------------------\n")
+        count_local_emails(args.path)
+        raise SystemExit(0)
+
+    IMAP_SERVER = args.host
+    USERNAME = args.user
+    PASSWORD = args.password
+
+    # Acquire OAuth2 token if configured
+    oauth2_token = None
+    oauth2_provider = None
+    if args.client_id:
+        oauth2_token, oauth2_provider = imap_oauth2.acquire_token(
+            IMAP_SERVER, args.client_id, USERNAME, args.client_secret
+        )
+
+    print("\n--- Configuration Summary ---")
+    print(f"Host            : {IMAP_SERVER}")
+    print(f"User            : {USERNAME}")
+    print(f"Auth Method     : {imap_oauth2.auth_description(oauth2_provider)}")
+    print("-----------------------------\n")
+
+    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_oauth2.py

@@ -1,258 +0,0 @@
-"""
-IMAP OAuth2 Authentication
-
-OAuth2 token acquisition and refresh for Microsoft and Google IMAP providers.
-Supports device code flow (Microsoft) and installed app flow (Google),
-with in-memory caching for silent token refresh.
-
-Notes:
-- Microsoft OAuth2 requires the `msal` package and uses device code flow.
-- Google OAuth2 requires the `google-auth-oauthlib` package and uses an installed-app flow
-    (opens a browser and listens on a local HTTP redirect).
-- Provider is auto-detected from the IMAP host string.
-"""
-
-import http.client
-import json
-import re
-import ssl
-import sys
-import threading
-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
-_token_refresh_lock = threading.Lock()
-
-
-def _fetch_json_https(host, path, timeout=10):
-    if not host or any(ch in host for ch in "\r\n"):
-        raise ValueError("Invalid host")
-    if not path.startswith("/"):
-        path = f"/{path}"
-
-    context = ssl.create_default_context()
-    conn = http.client.HTTPSConnection(host, timeout=timeout, context=context)
-    try:
-        conn.request("GET", path, headers={"Accept": "application/json"})
-        response = conn.getresponse()
-        body = response.read()
-    finally:
-        conn.close()
-
-    if response.status != 200:
-        raise RuntimeError(f"Unexpected HTTP status {response.status}")
-    return json.loads(body.decode("utf-8"))
-
-
-def detect_oauth2_provider(host):
-    """
-    Detects the OAuth2 provider from the IMAP host.
-    Returns "microsoft", "google", or None if unrecognized.
-    """
-    host_lower = host.lower()
-    if "outlook" in host_lower or "office365" in host_lower or "microsoft" in host_lower:
-        return "microsoft"
-    if "gmail" in host_lower or "google" in host_lower:
-        return "google"
-    return None
-
-
-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).
-    Returns the tenant ID string or None if discovery fails.
-    """
-    domain = email.split("@")[-1].strip()
-    if not domain:
-        print("Error: Could not discover Microsoft tenant: missing email domain")
-        return None
-
-    domain_quoted = urllib.parse.quote(domain, safe=".-")
-    path = f"/{domain_quoted}/.well-known/openid-configuration"
-
-    try:
-        data = _fetch_json_https("login.microsoftonline.com", path, timeout=10)
-    except (OSError, http.client.HTTPException, json.JSONDecodeError, RuntimeError, ValueError) as e:
-        print(f"Error: Could not discover Microsoft tenant for domain '{domain}': {e}")
-        return None
-
-    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)
-
-    print(f"Error: Could not extract tenant ID from issuer: {issuer}")
-    return None
-
-
-def acquire_microsoft_oauth2_token(client_id, email):
-    """
-    Acquires a Microsoft OAuth2 access token using the MSAL device code flow.
-    Auto-discovers tenant ID from the email domain.
-    Requires the 'msal' package: pip install msal
-
-    On subsequent calls, silently refreshes the token using the cached MSAL app
-    (which holds the refresh token in its in-memory cache).
-    """
-    tenant_id = discover_microsoft_tenant(email)
-    if not tenant_id:
-        return None
-
-    try:
-        import msal
-    except ImportError:
-        print("Error: 'msal' package is required for Microsoft OAuth2. Install it with: pip install msal")
-        sys.exit(1)
-
-    authority = f"https://login.microsoftonline.com/{tenant_id}"
-    scopes = ["https://outlook.office365.com/IMAP.AccessAsUser.All"]
-
-    # Reuse cached MSAL app so acquire_token_silent can access refresh tokens
-    cache_key = (client_id, tenant_id)
-    if cache_key in _msal_app_cache:
-        app = _msal_app_cache[cache_key]
-    else:
-        print(f"Discovered Microsoft tenant: {tenant_id}")
-        app = msal.PublicClientApplication(client_id, authority=authority)
-        _msal_app_cache[cache_key] = app
-
-    # Try cached/refreshed token first (handles refresh tokens automatically)
-    accounts = app.get_accounts()
-    if accounts:
-        result = app.acquire_token_silent(scopes, account=accounts[0])
-        if result and "access_token" in result:
-            return result["access_token"]
-
-    # Fall back to device code flow (first call or if refresh fails)
-    flow = app.initiate_device_flow(scopes=scopes)
-    if "user_code" not in flow:
-        print(f"Error: Could not initiate device flow: {flow.get('error_description', 'Unknown error')}")
-        return None
-
-    print(flow["message"])
-    result = app.acquire_token_by_device_flow(flow)
-
-    if "access_token" in result:
-        return result["access_token"]
-
-    print(f"Error: Could not acquire token: {result.get('error_description', 'Unknown error')}")
-    return None
-
-
-def acquire_google_oauth2_token(client_id, client_secret):
-    """
-    Acquires a Google OAuth2 access token using the installed app flow.
-    Opens a browser for user consent and runs a local HTTP server for the redirect.
-    Requires the 'google-auth-oauthlib' package: pip install google-auth-oauthlib
-
-    On subsequent calls, silently refreshes the token using the cached credentials
-    object (which holds the refresh token). No browser interaction needed for refresh.
-    """
-    # Try refreshing cached credentials first (no browser needed)
-    cache_key = (client_id, client_secret)
-    if cache_key in _google_creds_cache:
-        creds = _google_creds_cache[cache_key]
-        if creds and creds.refresh_token:
-            try:
-                import google.auth.transport.requests
-
-                creds.refresh(google.auth.transport.requests.Request())
-                if creds.token:
-                    return creds.token
-            except Exception:
-                pass  # Fall through to full auth flow
-
-    try:
-        from google_auth_oauthlib.flow import InstalledAppFlow
-    except ImportError:
-        print("Error: 'google-auth-oauthlib' package is required for Google OAuth2.")
-        print("Install it with: pip install google-auth-oauthlib")
-        sys.exit(1)
-
-    client_config = {
-        "installed": {
-            "client_id": client_id,
-            "client_secret": client_secret,
-            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
-            "token_uri": "https://oauth2.googleapis.com/token",
-            "redirect_uris": ["http://localhost"],
-        }
-    }
-
-    flow = InstalledAppFlow.from_client_config(client_config, scopes=["https://mail.google.com/"])
-
-    print("Opening browser for Google authentication...")
-    print("If the browser does not open, check the terminal for a URL to visit.")
-
-    credentials = flow.run_local_server(port=0)
-
-    if credentials and credentials.token:
-        _google_creds_cache[cache_key] = credentials
-        return credentials.token
-
-    print("Error: Could not acquire Google OAuth2 token.")
-    return None
-
-
-def acquire_oauth2_token_for_provider(provider, client_id, email, client_secret=None):
-    """
-    Acquires an OAuth2 token for the specified provider.
-
-    Args:
-        provider: "microsoft" or "google"
-        client_id: OAuth2 client ID
-        email: User's email address (used for Microsoft tenant discovery)
-        client_secret: Required for Google, not needed for Microsoft
-    """
-    if provider == "microsoft":
-        return acquire_microsoft_oauth2_token(client_id, email)
-    elif provider == "google":
-        if not client_secret:
-            print(
-                "Error: OAuth2 client secret is required for Google OAuth2. "
-                "Provide --oauth2-client-secret (or --src-oauth2-client-secret / --dest-oauth2-client-secret), "
-                "or set OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET / DEST_OAUTH2_CLIENT_SECRET."
-            )
-            return None
-        return acquire_google_oauth2_token(client_id, client_secret)
-    else:
-        print(f"Error: Unknown OAuth2 provider: {provider}")
-        return None
-
-
-def refresh_oauth2_token(conf, old_token):
-    """
-    Thread-safe OAuth2 token refresh using double-checked locking.
-
-    Multiple threads may detect an expired token simultaneously. This function
-    ensures only one thread performs the actual refresh. Other threads waiting
-    on the lock will see that conf["oauth2_token"] has already been updated and
-    skip the redundant refresh.
-
-    Args:
-        conf: Mutable dict with keys:
-            - host, user, password, oauth2_token: connection credentials
-            - oauth2: dict with provider, client_id, email, client_secret
-        old_token: The expired token that triggered this refresh (for comparison)
-
-    Returns:
-        The new token string, or None if refresh failed.
-    """
-    oauth2 = conf.get("oauth2")
-    if not oauth2:
-        return None
-
-    with _token_refresh_lock:
-        # Another thread may have already refreshed while we were waiting
-        if conf["oauth2_token"] != old_token:
-            return conf["oauth2_token"]
-
-        new_token = acquire_oauth2_token_for_provider(
-            oauth2["provider"], oauth2["client_id"], oauth2["email"], oauth2.get("client_secret")
-        )
-        if new_token:
-            conf["oauth2_token"] = new_token
-        return new_token

src/imap_restore.py

@@ -0,0 +1,976 @@
+"""
+IMAP Email Restore Script
+
+Restores emails from a local backup to an IMAP account.
+Reads .eml files from a local directory and uploads them to the destination server.
+
+Features:
+- Folder Restoration: Recreates the folder structure from the backup.
+- Gmail Labels Restoration: Uses labels_manifest.json to apply Gmail labels.
+- Incremental Restore: Skips emails that already exist (based on Message-ID).
+- Parallel Processing: Uses multithreading for fast uploads.
+- Date Preservation: Restores emails with their original dates.
+
+Configuration (Environment Variables):
+    DEST_IMAP_HOST, DEST_IMAP_USERNAME: Destination credentials.
+    DEST_IMAP_PASSWORD: Destination password (or App Password).
+
+    OAuth2 (Optional - instead of password):
+    DEST_OAUTH2_CLIENT_ID: OAuth2 Client ID
+    DEST_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
+
+  BACKUP_LOCAL_PATH: Source local directory containing the backup.
+  MAX_WORKERS: Number of concurrent threads (default: 4).
+  BATCH_SIZE: Number of emails to process per batch (default: 10).
+  APPLY_LABELS: Set to "true" to apply Gmail labels from manifest. Default is "false".
+  APPLY_FLAGS: Set to "true" to apply IMAP flags from manifest. Default is "false".
+  GMAIL_MODE: Set to "true" for Gmail restore mode. Default is "false".
+  DEST_DELETE: Set to "true" to delete emails from destination not found in local backup.
+              Default is "false".
+
+Usage:
+    python3 imap_restore.py \
+        --src-path "./my_backup" \
+        --dest-host "imap.gmail.com" \
+        --dest-user "you@gmail.com" \
+        --dest-pass "your-app-password"
+
+Gmail Labels Restoration:
+    python3 imap_restore.py \
+        --src-path "./gmail_backup" \
+        --dest-host "imap.gmail.com" \
+        --dest-user "you@gmail.com" \
+        --dest-pass "your-app-password" \
+        --apply-labels
+  This uploads emails and applies labels from labels_manifest.json to recreate
+  the original Gmail label structure.
+"""
+
+import argparse
+import concurrent.futures
+import os
+import sys
+import threading
+import time
+from email import policy
+from email.parser import BytesParser
+from email.utils import parsedate_to_datetime
+from enum import Enum
+from typing import Optional
+
+from auth import imap_oauth2
+from core import imap_session
+from providers import provider_gmail
+from utils import imap_common, restore_cache
+
+
+class UploadResult(Enum):
+    """Result of an email upload operation."""
+
+    SUCCESS = "success"
+    ALREADY_EXISTS = "already_exists"
+    FAILURE = "failure"
+
+
+# Defaults
+MAX_WORKERS = 4  # Lower default for restore to avoid rate limits
+BATCH_SIZE = 10
+
+# Thread-local storage
+thread_local = threading.local()
+safe_print = imap_common.safe_print
+
+
+def get_flags_from_manifest(manifest, message_id):
+    """
+    Extract IMAP flags string from manifest entry.
+    Returns flags string like "\\Seen \\Flagged" or None.
+    """
+    if not message_id or message_id not in manifest:
+        return None
+
+    entry = manifest[message_id]
+
+    if isinstance(entry, dict) and "flags" in entry:
+        flags = entry.get("flags", [])
+        if flags:
+            return " ".join(flags)
+
+    return None
+
+
+def parse_eml_file(file_path):
+    """
+    Parse an .eml file and extract metadata.
+    Returns (message_id, date_str, raw_content, subject) or (None, None, None, None) on error.
+    """
+    try:
+        with open(file_path, "rb") as f:
+            raw_content = f.read()
+
+        # Use compat32 to preserve raw headers with continuation lines
+        parser = BytesParser(policy=policy.compat32)
+        msg = parser.parsebytes(raw_content, headersonly=True)
+
+        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
+        date_str = None
+        if date_header:
+            try:
+                dt = parsedate_to_datetime(date_header)
+                # Format: "DD-Mon-YYYY HH:MM:SS +ZZZZ"
+                date_str = dt.strftime('"%d-%b-%Y %H:%M:%S %z"')
+            except Exception:
+                pass
+
+        return message_id, date_str, raw_content, subject
+
+    except Exception as e:
+        safe_print(f"Error parsing {file_path}: {e}")
+        return None, None, None, None
+
+
+def get_eml_files(folder_path):
+    """
+    Get all .eml files in a folder.
+    Returns list of (file_path, filename) tuples.
+    """
+    eml_files = []
+    if not os.path.exists(folder_path):
+        return eml_files
+
+    try:
+        for filename in os.listdir(folder_path):
+            if filename.endswith(".eml"):
+                file_path = os.path.join(folder_path, filename)
+                eml_files.append((file_path, filename))
+    except Exception as e:
+        safe_print(f"Error listing {folder_path}: {e}")
+
+    return eml_files
+
+
+def upload_email(dest, folder_name, raw_content, date_str, flags=None):
+    """
+    Upload a single email to the destination folder.
+    Returns UploadResult enum: SUCCESS or FAILURE.
+
+    Callers are expected to check for duplicates via load_folder_msg_ids()
+    before calling this function, and to ensure the folder exists.
+
+    Args:
+        flags: Optional string of IMAP flags like "\\Seen" for read emails.
+    """
+    try:
+        success = imap_common.append_email(
+            dest,
+            folder_name,
+            raw_content,
+            date_str,
+            flags,
+            ensure_folder=False,
+        )
+        return UploadResult.SUCCESS if success else UploadResult.FAILURE
+
+    except Exception as e:
+        safe_print(f"Error uploading to {folder_name}: {e}")
+        return UploadResult.FAILURE
+
+
+def get_labels_from_manifest(manifest, message_id):
+    """
+    Get the list of labels for a message from the manifest.
+    Returns list of label strings or empty list.
+    """
+    if not message_id or message_id not in manifest:
+        return []
+
+    entry = manifest[message_id]
+    if isinstance(entry, dict):
+        return entry.get("labels", [])
+    elif isinstance(entry, list):
+        return entry  # Old format: list of labels
+    return []
+
+
+def process_restore_batch(
+    eml_files,
+    folder_name,
+    dest_conf,
+    manifest,
+    apply_labels,
+    apply_flags,
+    full_restore: bool = False,
+    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None,
+    existing_dest_msg_ids_lock: Optional[threading.Lock] = None,
+    progress_cache_data: Optional[dict] = None,
+    progress_cache_lock: Optional[threading.Lock] = None,
+    progress_cache_path: Optional[str] = None,
+    dest_host: Optional[str] = None,
+    dest_user: Optional[str] = None,
+):
+    """
+    Process a batch of .eml files for restoration.
+
+    Args:
+        folder_name: Target folder, or "__GMAIL_MODE__" for per-email folder selection
+        manifest: Combined manifest with labels and/or flags
+        apply_labels: Whether to apply Gmail labels from manifest
+        apply_flags: Whether to apply IMAP flags from manifest
+    """
+    dest = imap_session.get_thread_connection(thread_local, "dest", dest_conf)
+    if not dest:
+        safe_print("Error: Could not establish connection for batch.")
+        return
+
+    gmail_mode = folder_name == "__GMAIL_MODE__"
+
+    for file_path, filename in eml_files:
+        # Proactively refresh token if needed
+        dest = imap_session.get_thread_connection(thread_local, "dest", 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  # No content, skip to next file
+
+            size = len(raw_content)
+            size_str = f"{size / 1024:.1f}KB"
+
+            # Truncate subject for display
+            display_subject = (subject[:40] + "...") if len(subject) > 40 else subject
+
+            # Get flags from manifest if apply_flags is enabled
+            flags = None
+            if apply_flags:
+                flags = get_flags_from_manifest(manifest, message_id)
+
+            # Get labels for this message
+            labels = get_labels_from_manifest(manifest, message_id) if apply_labels else []
+
+            # Determine target folder and remaining labels
+            if gmail_mode:
+                target_folder, remaining_labels = provider_gmail.resolve_target(labels)
+            else:
+                target_folder = folder_name
+                remaining_labels = labels
+
+            existing_dest_msg_ids = imap_common.load_folder_msg_ids(
+                dest,
+                target_folder,
+                existing_dest_msg_ids_by_folder,
+                existing_dest_msg_ids_lock,
+                progress_cache_data,
+                progress_cache_lock,
+                dest_host,
+                dest_user,
+            )
+
+            # Check if email already exists on destination using pre-fetched set
+            email_already_on_dest = (
+                message_id and existing_dest_msg_ids is not None and message_id in existing_dest_msg_ids
+            )
+
+            # Incremental default: skip entirely if already present
+            if email_already_on_dest and not full_restore:
+                safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
+                continue  # Skip to next file
+
+            if email_already_on_dest:
+                # Full restore: treat as existing for flag sync
+                upload_result = UploadResult.ALREADY_EXISTS
+            else:
+                # Upload — skip per-message SEARCH since we have a pre-fetched set
+                upload_result = upload_email(
+                    dest,
+                    target_folder,
+                    raw_content,
+                    date_str,
+                    flags,
+                )
+
+            # Only record progress when upload succeeds or email already exists.
+            # Failed uploads should not be marked as processed to allow retry on next run.
+            if upload_result in (UploadResult.SUCCESS, UploadResult.ALREADY_EXISTS):
+                restore_cache.record_progress(
+                    message_id=message_id,
+                    folder_name=target_folder,
+                    existing_dest_msg_ids=existing_dest_msg_ids,
+                    existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
+                    progress_cache_path=progress_cache_path,
+                    progress_cache_data=progress_cache_data,
+                    progress_cache_lock=progress_cache_lock,
+                    dest_host=dest_host,
+                    dest_user=dest_user,
+                    log_fn=safe_print,
+                )
+
+            if upload_result == UploadResult.ALREADY_EXISTS:
+                safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}")
+                # Full restore preserves legacy behavior: sync flags on existing email if requested
+                if full_restore and apply_flags and flags and message_id:
+                    imap_common.sync_flags_on_existing(dest, target_folder, message_id, flags, size)
+            elif upload_result == UploadResult.SUCCESS:
+                safe_print(f"[{target_folder}] UPLOADED      | {size_str:<8} | {display_subject}")
+                # Show applied flags in same style as labels
+                if flags:
+                    for flag in flags.split():
+                        safe_print(f"  -> Applied flag: {flag}")
+            else:  # FAILURE
+                safe_print(f"[{target_folder}] FAILED        | {size_str:<8} | {display_subject}")
+
+            # Apply remaining Gmail labels:
+            # - Full restore: apply/sync labels even for existing emails
+            # - Incremental (default): apply labels only for newly uploaded emails
+            if apply_labels and remaining_labels and (upload_result == UploadResult.SUCCESS or full_restore):
+                for label in remaining_labels:
+                    label_folder = provider_gmail.label_to_folder(label)
+
+                    # Skip if this is the same as the target folder
+                    if label_folder == target_folder:
+                        continue
+
+                    # Skip system folders we can't upload to
+                    if label_folder in (
+                        provider_gmail.GMAIL_ALL_MAIL,
+                        provider_gmail.GMAIL_SPAM,
+                        provider_gmail.GMAIL_TRASH,
+                    ):
+                        continue
+
+                    try:
+                        # Get or fetch Message-IDs for label folder (one-time server fetch per folder)
+                        label_folder_msg_ids = imap_common.load_folder_msg_ids(
+                            dest,
+                            label_folder,
+                            existing_dest_msg_ids_by_folder,
+                            existing_dest_msg_ids_lock,
+                            progress_cache_data,
+                            progress_cache_lock,
+                            dest_host,
+                            dest_user,
+                        )
+
+                        # Check duplicate using pre-fetched set instead of per-message SEARCH
+                        label_already_exists = (
+                            label_folder_msg_ids is not None and message_id and message_id in label_folder_msg_ids
+                        )
+
+                        if not label_already_exists:
+                            append_success = imap_common.append_email(
+                                dest,
+                                label_folder,
+                                raw_content,
+                                date_str,
+                                flags,
+                                ensure_folder=False,
+                            )
+                            if append_success:
+                                restore_cache.record_progress(
+                                    message_id=message_id,
+                                    folder_name=label_folder,
+                                    existing_dest_msg_ids=label_folder_msg_ids,
+                                    existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
+                                    progress_cache_path=progress_cache_path,
+                                    progress_cache_data=progress_cache_data,
+                                    progress_cache_lock=progress_cache_lock,
+                                    dest_host=dest_host,
+                                    dest_user=dest_user,
+                                    log_fn=safe_print,
+                                )
+                                safe_print(f"  -> Applied label: {label}")
+                            else:
+                                safe_print(f"  -> Failed to apply label {label} (will retry on next restore)")
+                        # If email exists in this label folder, sync flags (full restore only)
+                        elif full_restore and apply_flags and flags:
+                            imap_common.sync_flags_on_existing(dest, label_folder, message_id, flags, size)
+                    except Exception as e:
+                        if imap_oauth2.is_auth_error(e):
+                            raise
+                        safe_print(f"  -> Error applying label {label}: {e}")
+
+        except Exception as e:
+            safe_print(f"Error processing {filename}: {e}")
+
+
+def get_local_message_ids(local_folder_path):
+    """
+    Get a set of Message-IDs from all .eml files in a local folder.
+    Used for destination deletion sync.
+    """
+    message_ids = set()
+    eml_files = get_eml_files(local_folder_path)
+    for file_path, _filename in eml_files:
+        message_id, _, _, _ = parse_eml_file(file_path)
+        if message_id:
+            message_ids.add(message_id)
+    return message_ids
+
+
+def pre_filter_eml_files(eml_files, dest_msg_ids):
+    """Filter out .eml files whose Message-IDs already exist in the destination."""
+    safe_print("Pre-filtering duplicates...")
+    files_to_restore = []
+    skipped = 0
+    for file_path, filename in eml_files:
+        msg_id = imap_common.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.")
+    return files_to_restore
+
+
+def restore_folder(
+    folder_name,
+    local_folder_path,
+    dest_conf,
+    manifest,
+    apply_labels,
+    apply_flags,
+    dest_delete=False,
+    full_restore: bool = False,
+    cache_root: Optional[str] = None,
+    progress_cache_file: Optional[str] = None,
+    progress_cache_data: Optional[dict] = None,
+    progress_cache_lock: Optional[threading.Lock] = None,
+):
+    """
+    Restore all emails from a local folder to the destination IMAP server.
+    """
+    safe_print(f"--- Restoring Folder: {folder_name} ---")
+
+    eml_files = get_eml_files(local_folder_path)
+    if not eml_files:
+        safe_print(f"No .eml files found in {folder_name}")
+        # Even if empty, check for orphans to delete
+        if dest_delete:
+            dest = imap_common.get_imap_connection_from_conf(dest_conf)
+            if dest:
+                imap_common.delete_orphan_emails(dest, folder_name, set())
+                dest.logout()
+        return
+
+    safe_print(f"Found {len(eml_files)} emails to restore.")
+
+    cache_root = cache_root or local_folder_path
+    cache_path = progress_cache_file
+    cache_data = progress_cache_data
+    cache_lock = progress_cache_lock
+    if cache_path is None or cache_data is None or cache_lock is None:
+        cache_path, cache_data, cache_lock = imap_common.load_progress_cache(
+            cache_root,
+            dest_conf["host"],
+            dest_conf["user"],
+            log_fn=safe_print,
+        )
+
+    # Incremental mode uses cached Message-IDs to skip already-processed emails.
+    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {folder_name: set()}
+    existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
+    try:
+        existing_dest_msg_ids_by_folder[folder_name] = restore_cache.get_cached_message_ids(
+            cache_data,
+            cache_lock,
+            dest_conf["host"],
+            dest_conf["user"],
+            folder_name,
+        )
+        safe_print(f"Cache has {len(existing_dest_msg_ids_by_folder[folder_name])} Message-IDs for this folder.")
+    except Exception as e:
+        # Fall back to an empty cache for this folder if reading cached Message-IDs fails.
+        safe_print(f"Warning: Failed to load cached Message-IDs for folder '{folder_name}': {e}")
+        existing_dest_msg_ids_by_folder[folder_name] = set()
+
+    # If dest_delete enabled, get local Message-IDs for comparison
+    local_msg_ids = None
+    if dest_delete:
+        safe_print("Building local Message-ID index for sync...")
+        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.")
+
+        # Update destination Message-IDs with what we processed
+        if not full_restore:
+            with existing_dest_msg_ids_lock:
+                existing_dest_msg_ids_by_folder.setdefault(folder_name, set()).update(dest_msg_ids)
+        else:
+            existing_dest_msg_ids_by_folder[folder_name] = dest_msg_ids
+        dest_msg_ids = existing_dest_msg_ids_by_folder[folder_name]
+
+        # Pre-filter files to skip duplicates (skip in full_restore: need to process all for flag/label sync)
+        if dest_msg_ids and not full_restore:
+            files_to_restore = pre_filter_eml_files(eml_files, dest_msg_ids)
+
+    if not files_to_restore:
+        safe_print("No new emails to restore.")
+        if dest_delete and local_msg_ids is not None:
+            safe_print("Syncing destination: removing emails not in local backup...")
+            dest = imap_session.ensure_connection(None, dest_conf)
+            if dest:
+                imap_common.delete_orphan_emails(dest, folder_name, local_msg_ids)
+                dest.logout()
+
+        restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True)
+        return
+
+    safe_print(f"Starting parallel restore of {len(files_to_restore)} emails...")
+
+    # Create batches
+    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 = []
+        for batch in batches:
+            futures.append(
+                executor.submit(
+                    process_restore_batch,
+                    batch,
+                    folder_name,
+                    dest_conf,
+                    manifest,
+                    apply_labels,
+                    apply_flags,
+                    full_restore,
+                    existing_dest_msg_ids_by_folder,
+                    existing_dest_msg_ids_lock,
+                    cache_data,
+                    cache_lock,
+                    cache_path,
+                    dest_conf.get("host"),
+                    dest_conf.get("user"),
+                )
+            )
+
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                future.result()
+            except Exception as e:
+                safe_print(f"Batch error: {e}")
+
+    # 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 = imap_session.ensure_connection(None, dest_conf)
+        if dest:
+            imap_common.delete_orphan_emails(dest, folder_name, local_msg_ids)
+            dest.logout()
+
+    # Force-flush progress cache at end.
+    restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True)
+
+
+def restore_gmail_with_labels(
+    local_path,
+    dest_conf,
+    manifest,
+    apply_flags,
+    full_restore: bool = False,
+    progress_cache_file: Optional[str] = None,
+    progress_cache_data: Optional[dict] = None,
+    progress_cache_lock: Optional[threading.Lock] = None,
+):
+    """
+    Special restoration mode for Gmail: Upload emails to their first label folder
+    and then apply additional labels from the manifest.
+
+    This avoids putting all emails in INBOX - emails only appear in INBOX
+    if they originally had the INBOX label.
+    """
+    # Find the All Mail folder in the backup
+    all_mail_path = os.path.join(local_path, "[Gmail]", "All Mail")
+    if not os.path.exists(all_mail_path):
+        # Try without subfolder structure
+        all_mail_path = local_path
+
+    safe_print("--- Gmail Restore with Labels ---")
+    safe_print(f"Source path: {all_mail_path}")
+    safe_print(f"Entries in manifest: {len(manifest)}")
+
+    eml_files = get_eml_files(all_mail_path)
+    if not eml_files:
+        safe_print("No .eml files found for restoration.")
+        return
+
+    safe_print(f"Found {len(eml_files)} emails to restore.\n")
+
+    # Process in batches
+    total = len(eml_files)
+    start_time = time.time()
+
+    batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)]
+
+    cache_path = progress_cache_file
+    if cache_path is None or progress_cache_data is None or progress_cache_lock is None:
+        cache_path, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache(
+            local_path,
+            dest_conf["host"],
+            dest_conf["user"],
+            log_fn=safe_print,
+        )
+    safe_print(
+        "Cache will be populated as restore runs (no up-front destination indexing). "
+        "First run may still do per-message duplicate checks; subsequent runs will skip quickly."
+    )
+
+    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {}  # lazily loaded per folder
+    existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+        futures = []
+        for batch in batches:
+            # For Gmail, we use a special folder marker to indicate gmail-mode
+            # The process_restore_batch will determine the target folder per-email
+            # based on the manifest labels
+            futures.append(
+                executor.submit(
+                    process_restore_batch,
+                    batch,
+                    "__GMAIL_MODE__",  # Special marker - target determined per-email from manifest
+                    dest_conf,
+                    manifest,
+                    True,  # apply_labels
+                    apply_flags,  # apply_flags
+                    full_restore,
+                    existing_dest_msg_ids_by_folder,
+                    existing_dest_msg_ids_lock,
+                    progress_cache_data,
+                    progress_cache_lock,
+                    cache_path,
+                    dest_conf.get("host"),
+                    dest_conf.get("user"),
+                )
+            )
+
+        completed = 0
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                future.result()
+                completed += 1
+                elapsed = time.time() - start_time
+                progress = (completed / len(batches)) * 100
+                safe_print(f"Progress: {progress:.1f}% ({elapsed:.0f}s elapsed)")
+            except Exception as e:
+                safe_print(f"Batch error: {e}")
+
+    elapsed = time.time() - start_time
+    safe_print(f"\nRestore completed in {elapsed:.1f}s")
+
+    # Force-flush progress cache at end.
+    restore_cache.maybe_save_dest_index_cache(cache_path, progress_cache_data, progress_cache_lock, force=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Restore IMAP emails from local .eml files.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+
+    # Source (Local Path)
+    env_path = os.getenv("BACKUP_LOCAL_PATH")
+    parser.add_argument(
+        "--src-path",
+        default=env_path,
+        required=not bool(env_path),
+        help="Local source path containing backup (or BACKUP_LOCAL_PATH)",
+    )
+
+    # Destination
+    default_dest_host = os.getenv("DEST_IMAP_HOST")
+    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
+    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
+    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--dest-host",
+        default=default_dest_host,
+        required=not bool(default_dest_host),
+        help="Destination IMAP Server (or DEST_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--dest-user",
+        default=default_dest_user,
+        required=not bool(default_dest_user),
+        help="Destination Username (or DEST_IMAP_USERNAME)",
+    )
+
+    dest_auth_required = not bool(default_dest_pass or default_dest_client_id)
+    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
+    dest_auth.add_argument(
+        "--dest-pass",
+        default=default_dest_pass,
+        help="Destination Password (or DEST_IMAP_PASSWORD)",
+    )
+    dest_auth.add_argument(
+        "--dest-oauth2-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
+    )
+    dest_auth.add_argument(
+        "--dest-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--dest-oauth2-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--dest-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help=argparse.SUPPRESS,
+    )
+
+    # Config
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=int(os.getenv("MAX_WORKERS", 4)),
+        help="Thread count (default: 4)",
+    )
+    parser.add_argument(
+        "--batch",
+        type=int,
+        default=int(os.getenv("BATCH_SIZE", 10)),
+        help="Emails per batch",
+    )
+
+    # Gmail Labels
+    env_apply_labels = os.getenv("APPLY_LABELS", "false").lower() == "true"
+    parser.add_argument(
+        "--apply-labels",
+        action="store_true",
+        default=env_apply_labels,
+        help="Apply Gmail labels from labels_manifest.json",
+    )
+    env_apply_flags = os.getenv("APPLY_FLAGS", "false").lower() == "true"
+    parser.add_argument(
+        "--apply-flags",
+        action="store_true",
+        default=env_apply_flags,
+        help="Apply IMAP flags (read/starred/answered/draft) from manifest",
+    )
+    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
+    parser.add_argument(
+        "--gmail-mode",
+        action="store_true",
+        default=env_gmail_mode,
+        help="Gmail restore mode: Upload to INBOX and apply labels + flags from manifest",
+    )
+
+    env_full_restore = os.getenv("FULL_RESTORE", "false").lower() == "true"
+    parser.add_argument(
+        "--full-restore",
+        action="store_true",
+        default=env_full_restore,
+        help="Force full restore (legacy): process all emails and sync labels/flags for already-present messages.",
+    )
+
+    # Sync mode: delete from dest emails not in local backup
+    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
+    parser.add_argument(
+        "--dest-delete",
+        action="store_true",
+        default=env_dest_delete,
+        help="Delete emails from destination that don't exist in local backup (sync mode)",
+    )
+
+    # Optional folder filter
+    parser.add_argument("folder", nargs="?", help="Specific folder to restore")
+
+    args = parser.parse_args()
+
+    global MAX_WORKERS, BATCH_SIZE
+    MAX_WORKERS = args.workers
+    BATCH_SIZE = args.batch
+
+    # Build connection config (acquires OAuth2 token if configured)
+    dest_conf = imap_session.build_imap_conf(
+        args.dest_host, args.dest_user, args.dest_pass, args.dest_client_id, args.dest_client_secret, "destination"
+    )
+
+    # Expand path
+    local_path = os.path.expanduser(args.src_path)
+    if not os.path.exists(local_path):
+        print(f"Error: Source path does not exist: {local_path}")
+        sys.exit(1)
+
+    # Load manifest(s) if needed
+    manifest = {}
+    apply_labels = args.apply_labels or args.gmail_mode
+    apply_flags = args.apply_flags or args.gmail_mode
+
+    if apply_labels or apply_flags:
+        # Try to load labels manifest first (contains both labels and flags for Gmail backups)
+        manifest = imap_common.load_manifest(local_path, "labels_manifest.json")
+
+        # If no labels manifest, try flags-only manifest
+        if not manifest and apply_flags:
+            manifest = imap_common.load_manifest(local_path, "flags_manifest.json")
+
+        if not manifest:
+            print("Warning: No manifest found. Labels/flags will not be applied.")
+
+    progress_cache_file = None
+    progress_cache_data = None
+    progress_cache_lock = None
+    try:
+        progress_cache_file, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache(
+            local_path,
+            args.dest_host,
+            args.dest_user,
+            log_fn=safe_print,
+        )
+    except Exception as e:
+        safe_print(f"Warning: Failed to load progress cache: {e}")
+
+    print("\n--- Configuration Summary ---")
+    print(f"Source Path     : {local_path}")
+    print(f"Destination Host: {args.dest_host}")
+    print(f"Destination User: {args.dest_user}")
+    print(f"Destination Auth: {imap_oauth2.auth_description(dest_conf['oauth2'] and dest_conf['oauth2']['provider'])}")
+    print(f"Workers         : {args.workers}")
+    if args.gmail_mode:
+        print("Mode            : Gmail Restore with Labels + Flags")
+    elif args.folder:
+        print(f"Target Folder   : {args.folder}")
+    if apply_labels:
+        print(f"Apply Labels    : Yes ({len(manifest)} mappings)")
+    if apply_flags:
+        print("Apply Flags     : Yes (read/starred/answered/draft)")
+    if args.dest_delete:
+        print("Dest Delete     : Yes (remove orphans from destination)")
+    print(
+        f"Restore Mode    : {'Full (all emails)' if args.full_restore else 'Incremental (new emails only, use --full-restore to restore all)'}"
+    )
+    print("-----------------------------\n")
+
+    try:
+        # Test connection
+        dest = imap_common.get_imap_connection_from_conf(dest_conf)
+        if not dest:
+            print("Error: Could not connect to destination server.")
+            sys.exit(1)
+
+        if args.gmail_mode:
+            dest.logout()
+            # Special Gmail mode
+            restore_gmail_with_labels(
+                local_path,
+                dest_conf,
+                manifest,
+                apply_flags,
+                full_restore=args.full_restore,
+                progress_cache_file=progress_cache_file,
+                progress_cache_data=progress_cache_data,
+                progress_cache_lock=progress_cache_lock,
+            )
+            dest = None  # Connection handled by restore_gmail_with_labels
+        elif args.folder:
+            # Restore specific folder
+            folder_path = os.path.join(local_path, args.folder.replace("/", os.sep))
+            if not os.path.exists(folder_path):
+                print(f"Error: Folder not found: {folder_path}")
+                sys.exit(1)
+            restore_folder(
+                args.folder,
+                folder_path,
+                dest_conf,
+                manifest,
+                apply_labels,
+                apply_flags,
+                args.dest_delete,
+                full_restore=args.full_restore,
+                cache_root=local_path,
+                progress_cache_file=progress_cache_file,
+                progress_cache_data=progress_cache_data,
+                progress_cache_lock=progress_cache_lock,
+            )
+            dest.logout()
+        else:
+            # Restore all folders
+            folders = imap_common.get_backup_folders(local_path)
+            if not folders:
+                print("No backup folders found.")
+                sys.exit(1)
+
+            print(f"Found {len(folders)} folders to restore.\n")
+            for folder_name, folder_path in folders:
+                # Skip manifest files
+                if folder_name in ("labels_manifest.json", "flags_manifest.json"):
+                    continue
+
+                # Proactively refresh OAuth2 token and ensure connection is healthy between folders
+                dest = imap_session.ensure_connection(dest, dest_conf)
+                if not dest:
+                    print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
+                    sys.exit(1)
+
+                restore_folder(
+                    folder_name,
+                    folder_path,
+                    dest_conf,
+                    manifest,
+                    apply_labels,
+                    apply_flags,
+                    args.dest_delete,
+                    full_restore=args.full_restore,
+                    cache_root=local_path,
+                    progress_cache_file=progress_cache_file,
+                    progress_cache_data=progress_cache_data,
+                    progress_cache_lock=progress_cache_lock,
+                )
+
+            dest.logout()
+
+        print("\nRestore completed successfully.")
+
+    except KeyboardInterrupt:
+        raise
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)

src/providers/provider_exchange.py

@@ -0,0 +1,29 @@
+"""
+Exchange-Specific IMAP Utilities
+
+Constants and functions specific to Microsoft Exchange/Outlook IMAP implementation.
+"""
+
+# Exchange/Outlook folders that typically can't be backed up via IMAP
+# These often contain proprietary data that Exchange returns as error messages
+EXCHANGE_SKIP_FOLDERS = {
+    "Calendar",
+    "Contacts",
+    "Conversation History",
+    "Suggested Contacts",
+}
+
+
+def is_special_folder(folder_name):
+    """
+    Check if a folder is an Exchange system folder that should be skipped.
+
+    These folders often contain proprietary data that Exchange returns as error messages.
+
+    Args:
+        folder_name: The name of the folder to check
+
+    Returns:
+        True if the folder should be skipped, False otherwise
+    """
+    return folder_name in EXCHANGE_SKIP_FOLDERS

src/providers/provider_gmail.py

@@ -0,0 +1,130 @@
+"""
+Gmail-Specific IMAP Utilities
+
+Constants and functions specific to Gmail/Google Workspace IMAP implementation.
+"""
+
+from auth import imap_oauth2
+from utils import imap_common
+
+# Gmail system folders
+GMAIL_ALL_MAIL = "[Gmail]/All Mail"
+GMAIL_TRASH = "[Gmail]/Trash"
+GMAIL_SPAM = "[Gmail]/Spam"
+GMAIL_DRAFTS = "[Gmail]/Drafts"
+GMAIL_BIN = "[Gmail]/Bin"
+GMAIL_IMPORTANT = "[Gmail]/Important"
+GMAIL_SENT = "[Gmail]/Sent Mail"
+GMAIL_STARRED = "[Gmail]/Starred"
+
+GMAIL_SYSTEM_FOLDERS = {
+    GMAIL_ALL_MAIL,
+    GMAIL_SPAM,
+    GMAIL_TRASH,
+    GMAIL_DRAFTS,
+    GMAIL_BIN,
+    GMAIL_IMPORTANT,
+}
+
+
+def is_label_folder(folder_name):
+    """
+    Determines if a folder represents a Gmail label (user-created or system label
+    that should be preserved).
+    Excludes system folders like All Mail, Spam, Trash, Drafts.
+    """
+    # Exclude system folders that aren't really "labels"
+    if folder_name in GMAIL_SYSTEM_FOLDERS:
+        return False
+
+    # INBOX is a special case - it's a label in Gmail
+    if folder_name == imap_common.FOLDER_INBOX:
+        return True
+
+    # [Gmail]/Sent Mail and [Gmail]/Starred are labels worth preserving
+    if folder_name in (GMAIL_SENT, GMAIL_STARRED):
+        return True
+
+    # Any folder NOT under [Gmail]/ is a user label
+    if not folder_name.startswith("[Gmail]/"):
+        return True
+
+    return False
+
+
+def folder_to_label(folder_name):
+    """Convert an IMAP folder name to a Gmail label name (backup/restore compatible)."""
+    if folder_name == imap_common.FOLDER_INBOX:
+        return imap_common.FOLDER_INBOX
+    if folder_name.startswith("[Gmail]/"):
+        return folder_name.split("/", 1)[1]
+    return folder_name
+
+
+def label_to_folder(label):
+    """Convert a Gmail label name to an IMAP folder path (restore compatible)."""
+    if label == imap_common.FOLDER_INBOX:
+        return imap_common.FOLDER_INBOX
+    if label in ("Sent Mail", "Starred", "Drafts", "Important"):
+        return f"[Gmail]/{label}"
+    return label
+
+
+def resolve_target(labels):
+    """Pick the primary target folder and remaining labels for Gmail mode.
+
+    Returns:
+        Tuple of (target_folder, remaining_labels)
+    """
+    skip_folders = {GMAIL_ALL_MAIL, GMAIL_SPAM, GMAIL_TRASH}
+    target_folder = None
+    remaining_labels = []
+
+    for label in labels:
+        lf = label_to_folder(label)
+        if lf in skip_folders:
+            continue
+        if target_folder is None:
+            target_folder = lf
+        else:
+            remaining_labels.append(label)
+
+    if target_folder is None:
+        target_folder = imap_common.FOLDER_RESTORED_UNLABELED
+        remaining_labels = []
+
+    return target_folder, remaining_labels
+
+
+def build_gmail_label_index(src_conn, safe_print_func):
+    """
+    Build a mapping of Message-ID -> set(labels) by scanning label folders.
+
+    Args:
+        src_conn: IMAP connection to source account
+        safe_print_func: Function to use for printing (e.g., safe_print from the calling script)
+
+    Returns:
+        dict: Mapping of Message-ID -> set(label names)
+    """
+    folders = imap_common.list_selectable_folders(src_conn)
+    label_folders = [f for f in folders if is_label_folder(f)]
+
+    label_index = {}
+    total = len(label_folders)
+    for i, folder in enumerate(label_folders, start=1):
+        safe_print_func(f"[{i}/{total}] Scanning label folder for Message-IDs: {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:
+            # Re-raise auth errors so caller can handle reconnection
+            if imap_oauth2.is_auth_error(e):
+                raise
+            safe_print_func(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)
+
+    return label_index

src/utils/imap_common.py

@@ -0,0 +1,785 @@
+"""
+IMAP Common Utilities
+
+Shared functionality for IMAP migration, counting, and comparison scripts.
+"""
+
+from __future__ import annotations
+
+import imaplib
+import json
+import os
+import re
+import sys
+import threading
+import urllib.parse
+from email import policy
+from email.header import decode_header
+from email.parser import BytesParser
+from importlib.metadata import PackageNotFoundError, version
+
+from auth import imap_oauth2
+from utils import imap_compress, imap_retry, restore_cache
+
+
+def get_version() -> str:
+    """
+    Return the package version string.
+
+    Tries to get the version from the installed package metadata.
+    If not installed, falls back to reading pyproject.toml from the project root.
+    """
+    try:
+        return version("imap-migration-tools")
+    except PackageNotFoundError:
+        # Fallback: Try to read from pyproject.toml for local development
+        try:
+            # Assuming src/imap_common.py structure
+            root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            pyproject_path = os.path.join(root_dir, "pyproject.toml")
+
+            if os.path.exists(pyproject_path):
+                with open(pyproject_path, encoding="utf-8") as f:
+                    for line in f:
+                        # Looking for: version = "1.2.3"
+                        match = re.search(r'^version\s*=\s*"([^"]+)"', line.strip())
+                        if match:
+                            return match.group(1)
+        except Exception:
+            pass
+        return "0.0.0-dev"
+
+
+# Standard IMAP flags
+FLAG_SEEN = "\\Seen"
+FLAG_ANSWERED = "\\Answered"
+FLAG_FLAGGED = "\\Flagged"
+FLAG_DRAFT = "\\Draft"
+FLAG_DELETED = "\\Deleted"
+FLAG_DELETED_LITERAL = "(\\Deleted)"
+
+# Standard IMAP flags that can be preserved during migration
+# \Recent is session-specific and cannot be set by clients
+# \Deleted should not be preserved as it marks messages for removal
+PRESERVABLE_FLAGS = {FLAG_SEEN, FLAG_ANSWERED, FLAG_FLAGGED, FLAG_DRAFT}
+
+# IMAP Folder Constants
+FOLDER_INBOX = "INBOX"
+
+# Gmail-mode restore/migrate fallback folder for messages with no usable labels.
+# Keeping this as a normal folder (not [Gmail]/Drafts) avoids populating Drafts with non-drafts.
+FOLDER_RESTORED_UNLABELED = "Restored/Unlabeled"
+
+# IMAP Commands
+CMD_STORE = "store"
+CMD_SEARCH = "search"
+CMD_FETCH = "fetch"
+OP_ADD_FLAGS = "+FLAGS"
+
+_print_lock = threading.Lock()
+
+
+def safe_print(message: str) -> None:
+    """Thread-safe print with short thread names for logs."""
+    t_name = threading.current_thread().name
+    short_name = t_name.replace("ThreadPoolExecutor-", "T-").replace("MainThread", "MAIN")
+    with _print_lock:
+        print(f"[{short_name}] {message}")
+
+
+def _is_ignored_local_dir(dirname: str) -> bool:
+    return dirname.startswith(".") or dirname == "__pycache__"
+
+
+def list_local_folders(local_root: str) -> list[str]:
+    """List all folders under a local backup root in IMAP-style names."""
+    folders: set[str] = set()
+
+    for dirpath, dirnames, _filenames in os.walk(local_root):
+        dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)]
+
+        if os.path.abspath(dirpath) == os.path.abspath(local_root):
+            continue
+
+        rel = os.path.relpath(dirpath, local_root)
+        if rel == ".":
+            continue
+
+        parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)]
+        if not parts:
+            continue
+
+        folders.add("/".join(parts))
+
+    return sorted(folders)
+
+
+def get_local_email_count(local_root: str, folder_name: str) -> int | None:
+    """Return the count of .eml files in a local folder, or None if missing/unreadable."""
+    folder_path = os.path.join(local_root, *folder_name.split("/"))
+    if not os.path.isdir(folder_path):
+        return None
+
+    try:
+        count = 0
+        for filename in os.listdir(folder_path):
+            if not filename.endswith(".eml"):
+                continue
+            full_path = os.path.join(folder_path, filename)
+            if os.path.isfile(full_path):
+                count += 1
+        return count
+    except OSError:
+        return None
+
+
+def get_backup_folders(local_path: str) -> list[tuple[str, str]]:
+    """Scan the backup directory and return list of (folder_name, folder_path) tuples."""
+    folders: list[tuple[str, str]] = []
+
+    def scan_dir(path: str, prefix: str = "") -> None:
+        try:
+            for item in os.listdir(path):
+                item_path = os.path.join(path, item)
+                if os.path.isdir(item_path):
+                    # Check if this directory contains .eml files
+                    has_eml = any(
+                        f.endswith(".eml") for f in os.listdir(item_path) if os.path.isfile(os.path.join(item_path, f))
+                    )
+                    folder_name = f"{prefix}{item}" if prefix else item
+
+                    if has_eml:
+                        folders.append((folder_name, item_path))
+
+                    # Recurse into subdirectories
+                    scan_dir(item_path, f"{folder_name}/")
+        except Exception:
+            pass
+
+    scan_dir(local_path)
+    return folders
+
+
+def extract_message_id_from_eml(file_path: str, read_limit: int = 65536) -> str | None:
+    """Extract just the Message-ID from an .eml file efficiently."""
+    try:
+        with open(file_path, "rb") as f:
+            header_bytes = f.read(read_limit)
+
+        return extract_message_id(header_bytes)
+    except Exception:
+        return None
+
+
+def is_progress_cache_ready(cache_data: dict | None, cache_lock: threading.Lock | None) -> bool:
+    """Return True when cache data and lock are initialized."""
+    return cache_data is not None and cache_lock is not None
+
+
+def load_progress_cache(
+    cache_root: str,
+    dest_host: str,
+    dest_user: str,
+    *,
+    log_fn=None,
+) -> tuple[str, dict, threading.Lock]:
+    """Load or initialize a progress cache file for a destination."""
+    try:
+        os.makedirs(cache_root, exist_ok=True)
+    except Exception as exc:
+        if log_fn is not None:
+            log_fn(f"Warning: unable to create cache directory '{cache_root}': {exc}")
+
+    cache_path = restore_cache.get_dest_index_cache_path(cache_root, dest_host, dest_user)
+    cache_data = restore_cache.load_dest_index_cache(cache_path)
+    cache_lock = threading.Lock()
+    if log_fn is not None:
+        log_fn(f"Using progress cache: {cache_path}")
+    return cache_path, cache_data, cache_lock
+
+
+def ensure_folder_exists(imap_conn, folder_name: str) -> None:
+    """Best-effort create of a folder if it doesn't already exist.
+
+    IMAP servers typically return an error if the mailbox exists; this helper
+    intentionally ignores those errors.
+    """
+    try:
+        if folder_name and folder_name.upper() != FOLDER_INBOX:
+            imap_conn.create(f'"{folder_name}"')
+    except Exception:
+        # Folder may already exist, or server may restrict creation.
+        pass
+
+
+def append_email(
+    imap_conn,
+    folder_name: str,
+    raw_content: bytes,
+    date_str: str,
+    flags: str | None = None,
+    *,
+    ensure_folder: bool = True,
+) -> bool:
+    """Append an email message to a folder.
+
+    This is intentionally a thin wrapper around IMAP APPEND; callers can
+    perform duplicate checks or folder selection separately.
+
+    Args:
+        flags: Optional IMAP flags. If provided, they are normalized to a
+            parenthesized list before being passed to `imaplib.IMAP4.append`.
+        ensure_folder: If True, attempts to create the folder first (best-effort).
+    """
+    if ensure_folder:
+        ensure_folder_exists(imap_conn, folder_name)
+
+    normalized_flags = None
+    if flags:
+        stripped = str(flags).strip()
+        if stripped:
+            if stripped.startswith("(") and stripped.endswith(")"):
+                normalized_flags = stripped
+            else:
+                normalized_flags = f"({stripped})"
+
+    resp, data = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
+    if resp != "OK":
+        safe_print(f"APPEND failed for {folder_name}: {resp} {data}")
+    return resp == "OK"
+
+
+def verify_env_vars(vars_list):
+    """
+    Checks if all environment variables in the list are set.
+    Returns True if all are present, False otherwise.
+    Prints missing variables to stderr.
+    """
+    missing = [v for v in vars_list if not os.getenv(v)]
+    if missing:
+        print(f"Error: Missing environment variables: {', '.join(missing)}", file=sys.stderr)
+        return False
+    return True
+
+
+def get_imap_connection_from_conf(conf):
+    """
+    Establishes an IMAP connection using a conf dict.
+
+    conf dict structure:
+        {
+            "host": str,
+            "user": str,
+            "password": str or None,
+            "oauth2_token": str or None,
+            "oauth2": dict or None  # Contains provider, client_id, email, client_secret
+        }
+    """
+    return get_imap_connection(conf["host"], conf["user"], conf.get("password"), conf.get("oauth2_token"))
+
+
+def get_imap_connection(host, user, password=None, oauth2_token=None):
+    """
+    Establishes an SSL connection to the IMAP server and logs in.
+    Supports both basic auth (password) and OAuth 2.0 (XOAUTH2).
+    Returns the connection object or None if failed.
+    """
+    if not host or not user:
+        print(f"Error: Invalid credentials for {host}")
+        return None
+
+    if not password and not oauth2_token:
+        print(f"Error: Either password or oauth2_token is required for {host}")
+        return None
+
+    try:
+        use_ssl = True
+        resolved_host = host
+        port = None
+        if "://" in host:
+            parsed = urllib.parse.urlparse(host)
+            scheme = parsed.scheme.lower()
+            if not scheme or not parsed.hostname:
+                raise ValueError("Invalid IMAP host")
+            if scheme in {"imap", "tcp"}:
+                use_ssl = False
+            elif scheme in {"imaps", "imap+ssl", "imapssl", "ssl"}:
+                use_ssl = True
+            else:
+                raise ValueError(f"Unsupported IMAP scheme: {scheme}")
+            resolved_host = parsed.hostname
+            port = parsed.port
+
+        if use_ssl:
+            conn = imaplib.IMAP4_SSL(resolved_host, port) if port else imaplib.IMAP4_SSL(resolved_host)
+        else:
+            conn = imaplib.IMAP4(resolved_host, port) if port else imaplib.IMAP4(resolved_host)
+        if oauth2_token:
+            auth_string = f"user={user}\x01auth=Bearer {oauth2_token}\x01\x01"
+            conn.authenticate("XOAUTH2", lambda _: auth_string.encode())
+        else:
+            conn.login(user, password)
+        imap_compress.enable_compression(conn, log_fn=safe_print)
+        return imap_retry.ConnectionProxy(conn, log_fn=safe_print)
+    except Exception as e:
+        print(f"Connection error to {host}: {e}")
+        return None
+
+
+def ensure_connection(conn, host, user, password=None, oauth2_token=None):
+    """
+    Verifies an IMAP connection is still alive, reconnecting if necessary.
+    Returns the existing connection if healthy, or a new connection if it was broken.
+    Returns None if reconnection fails.
+    """
+    try:
+        if conn:
+            conn.noop()
+            return conn
+    except Exception:
+        # Connection is broken (network error, timeout, etc.) - fall through to reconnect
+        pass
+    return get_imap_connection(host, user, password, oauth2_token)
+
+
+def ensure_connection_from_conf(conn, conf):
+    """
+    Verifies an IMAP connection is still alive, reconnecting if necessary.
+    Uses a conf dict for connection parameters.
+    Returns the existing connection if healthy, or a new connection if it was broken.
+    Returns None if reconnection fails.
+    """
+    try:
+        if conn:
+            conn.noop()
+            return conn
+    except Exception:
+        # Connection is broken (network error, timeout, etc.) - fall through to reconnect
+        pass
+    return get_imap_connection_from_conf(conf)
+
+
+def normalize_folder_name(folder_info_str):
+    """
+    Parses the IMAP list response to extract the clean folder name.
+    Handles quoted names and flags.
+    """
+    if isinstance(folder_info_str, bytes):
+        folder_info_str = folder_info_str.decode("utf-8", errors="ignore")
+
+    # Regex to extract folder name: (flags) "delimiter" name
+    # Matches: (\HasNoChildren) "/" "INBOX"  OR  (\HasNoChildren) "/" Drafts
+    list_pattern = re.compile(r'\((?P<flags>.*?)\) "(?P<delimiter>.*)" "?(?P<name>.*)"?')
+    match = list_pattern.search(folder_info_str)
+    if match:
+        name = match.group("name")
+        # If the regex grabbed a trailing quote, strip it (though the regex tries to handle it)
+        return name.rstrip('"').strip()
+
+    # Fallback: take the last part
+    return folder_info_str.split()[-1].strip('"')
+
+
+def list_selectable_folders(imap_conn):
+    """
+    Lists all selectable folders (excluding \\Noselect) on the IMAP connection.
+    Returns a list of normalized folder name strings, or an empty list on failure.
+    """
+    try:
+        status, folders = imap_conn.list()
+        if status != "OK" or not folders:
+            return []
+    except Exception:
+        return []
+
+    result = []
+    for f in folders:
+        f_str = f.decode("utf-8", errors="ignore") if isinstance(f, bytes) else str(f)
+        if "\\Noselect" in f_str:
+            continue
+        result.append(normalize_folder_name(f))
+    return result
+
+
+def decode_mime_header(header_value):
+    """
+    Decodes MIME encoded headers (Subject, etc.) to a unicode string.
+    """
+    if not header_value:
+        return "(No Subject)"
+    try:
+        decoded_list = decode_header(header_value)
+        text_parts = []
+        for data, encoding in decoded_list:
+            if isinstance(data, bytes):
+                charset = encoding or "utf-8"
+                try:
+                    text_parts.append(data.decode(charset, errors="ignore"))
+                except LookupError:
+                    text_parts.append(data.decode("utf-8", errors="ignore"))
+            else:
+                text_parts.append(str(data))
+        return "".join(text_parts)
+    except Exception:
+        return str(header_value)
+
+
+def decode_message_id(msg_id):
+    """
+    Decodes a Message-ID header value by unfolding continuation lines.
+    Returns the stripped Message-ID string or None if empty.
+    """
+    if not msg_id:
+        return None
+    # Unfold any header continuation lines (CRLF/LF + whitespace)
+    return re.sub(r"\r?\n[ \t]+", " ", str(msg_id)).strip() or None
+
+
+def extract_message_id(header_data):
+    """
+    Extracts the Message-ID from header bytes or string using BytesParser.
+    Returns the stripped Message-ID string or None if not found.
+    """
+    if not header_data:
+        return None
+
+    try:
+        # Use compat32 to preserve raw header with continuation lines
+        # (policy.default's _MessageIDHeader parser truncates at newlines)
+        parser = BytesParser(policy=policy.compat32)
+        if isinstance(header_data, str):
+            header_data = header_data.encode("utf-8", errors="ignore")
+
+        msg = parser.parsebytes(header_data, headersonly=True)
+        return decode_message_id(msg.get("Message-ID"))
+    except Exception:
+        pass
+    return None
+
+
+def parse_message_id_from_bytes(raw_message):
+    """Parse Message-ID from RFC822 bytes.
+
+    Parses the full message bytes to extract the Message-ID header.
+    """
+    if not raw_message:
+        return None
+    try:
+        parser = BytesParser()
+        msg = parser.parsebytes(raw_message)
+        return decode_message_id(msg.get("Message-ID"))
+    except Exception:
+        return None
+
+
+def parse_message_id_and_subject_from_bytes(raw_message):
+    """Parse Message-ID and decoded Subject from RFC822 bytes.
+
+    Uses header-only parsing to avoid walking large message bodies.
+    Returns a tuple: (message_id, subject).
+    """
+    if not raw_message:
+        return None, "(No Subject)"
+
+    try:
+        # Use compat32 to preserve raw headers with continuation lines
+        parser = BytesParser(policy=policy.compat32)
+        email_obj = parser.parsebytes(raw_message, headersonly=True)
+        msg_id = decode_message_id(email_obj.get("Message-ID"))
+        raw_subject = email_obj.get("Subject")
+        subject = decode_mime_header(raw_subject) if raw_subject else "(No Subject)"
+        return msg_id, subject
+    except Exception:
+        return None, "(No Subject)"
+
+
+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 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()
+    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 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
+
+
+def sanitize_filename(filename):
+    """
+    Sanitizes a string to be safe for use as a filename.
+    Removes/replaces characters that are illegal in file systems.
+    Truncates to 250 chars.
+    """
+    if not filename:
+        return "untitled"
+    # Replace invalid characters with underscore
+    # Invalid: < > : " / \ | ? * and control chars
+    s = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", filename)
+    # Strip leading/trailing whitespaces/dots
+    s = s.strip().strip(".")
+    # Ensure not empty and not too long
+    return s[:250] if s else "untitled"
+
+
+def detect_trash_folder(imap_conn):
+    """
+    Attempts to identify the Trash folder in the account.
+    Returns the folder name (str) or None if not found.
+    Checks for common names and SPECIAL-USE attributes.
+    """
+    try:
+        status, folders = imap_conn.list()
+        if status != "OK":
+            return None
+    except Exception:
+        return None
+
+    trash_candidates = ["[Gmail]/Trash", "Trash", "Deleted Items", "Bin", "[Gmail]/Bin"]
+    detected_by_flag = None
+    all_folder_names = []
+
+    for f in folders:
+        if isinstance(f, bytes):
+            f_str = f.decode("utf-8", errors="ignore")
+        else:
+            f_str = str(f)
+
+        name = normalize_folder_name(f_str)
+        all_folder_names.append(name)
+
+        # Check for SPECIAL-USE flag \Trash
+        # The flag is usually inside parentheses like (\HasNoChildren \Trash)
+        if "\\Trash" in f_str or "\\Bin" in f_str:
+            detected_by_flag = name
+
+    if detected_by_flag:
+        return detected_by_flag
+
+    # Check candidates
+    for candidate in trash_candidates:
+        if candidate in all_folder_names:
+            return candidate
+
+    return None
+
+
+def sync_flags_on_existing(imap_conn, folder_name, message_id, flags, size):
+    """Sync flags on an existing email in the given folder.
+
+    Finds the email by Message-ID and adds any missing flags.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder containing the email
+        message_id: Message-ID header value
+        flags: Space-separated flags string like "\\Seen \\Flagged"
+        size: Email size for verification
+    """
+    if not flags or not message_id:
+        return
+
+    try:
+        imap_conn.select(f'"{folder_name}"')
+
+        clean_id = message_id.strip("<>").replace('"', '\\"')
+        typ, data = imap_conn.search(None, f'HEADER Message-ID "{clean_id}"')
+        if typ != "OK" or not data or not data[0]:
+            return
+
+        msg_num = data[0].split()[0]
+
+        flag_list = flags.split()
+        if not flag_list:
+            return
+
+        typ, flag_data = imap_conn.fetch(msg_num, "(FLAGS)")
+        if typ != "OK" or not flag_data:
+            return
+
+        current_flags = set()
+        for item in flag_data:
+            if isinstance(item, tuple) and item[0]:
+                resp_str = item[0].decode("utf-8", errors="ignore")
+                match = re.search(r"FLAGS\s+\((.*?)\)", resp_str)
+                if match:
+                    current_flags.update(match.group(1).split())
+
+        current_flags_lower = {f.lower() for f in current_flags}
+        flags_to_add = [f for f in flag_list if f.lower() not in current_flags_lower]
+
+        if flags_to_add:
+            flags_str = " ".join(flags_to_add)
+            typ, data = imap_conn.store(msg_num, "+FLAGS", f"({flags_str})")
+            if typ == "OK":
+                for flag in flags_to_add:
+                    safe_print(f"  -> Synced flag: {flag}")
+            else:
+                safe_print(f"STORE +FLAGS failed for {message_id} in {folder_name}: {typ} {data}")
+
+    except Exception as e:
+        safe_print(f"Error syncing flags for {message_id} in {folder_name}: {e}")
+
+
+def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msgid=None):
+    """Delete emails from a folder that don't exist in the source set.
+
+    Returns count of deleted emails.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to delete orphans from
+        source_msg_ids: Set of Message-IDs that should be kept
+        dest_uid_to_msgid: Optional pre-fetched dict of UID -> Message-ID.
+            If None, fetches from the server.
+    """
+    deleted_count = 0
+    try:
+        imap_conn.select(f'"{folder_name}"', readonly=False)
+
+        if dest_uid_to_msgid is None:
+            dest_uid_to_msgid = get_message_ids_in_folder(imap_conn)
+
+        uids_to_delete = []
+        for uid, msg_id in dest_uid_to_msgid.items():
+            if msg_id not in source_msg_ids:
+                uid_str = uid.decode() if isinstance(uid, bytes) else str(uid)
+                uids_to_delete.append(uid_str)
+
+        for uid in uids_to_delete:
+            try:
+                imap_conn.uid(CMD_STORE, uid, OP_ADD_FLAGS, FLAG_DELETED_LITERAL)
+                deleted_count += 1
+            except Exception as e:
+                safe_print(f"Warning: Failed to mark UID {uid} as deleted in folder {folder_name}: {e}")
+
+        if deleted_count > 0:
+            imap_conn.expunge()
+            safe_print(f"[{folder_name}] Deleted {deleted_count} orphan emails from destination")
+
+    except Exception as e:
+        safe_print(f"Error deleting orphans from {folder_name}: {e}")
+
+    return deleted_count
+
+
+def load_folder_msg_ids(
+    imap_conn,
+    folder_name,
+    msg_ids_by_folder,
+    msg_ids_lock,
+    progress_cache_data=None,
+    progress_cache_lock=None,
+    dest_host=None,
+    dest_user=None,
+):
+    """Load Message-IDs for a folder, fetching from server if not yet cached.
+
+    On first call for a folder, SELECTs the folder and fetches all Message-IDs
+    from the server (merged with progress cache). Subsequent calls return the
+    cached set without any IMAP operations.
+
+    Returns the set of Message-IDs, or None if tracking is disabled.
+    """
+    if msg_ids_by_folder is None or msg_ids_lock is None:
+        return None
+
+    with msg_ids_lock:
+        existing = msg_ids_by_folder.get(folder_name)
+
+    if existing is not None:
+        return existing
+
+    # Build from progress cache
+    built: set[str] = set()
+    if is_progress_cache_ready(progress_cache_data, progress_cache_lock) and dest_host and dest_user:
+        built = restore_cache.get_cached_message_ids(
+            progress_cache_data,
+            progress_cache_lock,
+            dest_host,
+            dest_user,
+            folder_name,
+        )
+
+    # Fetch from server for a comprehensive set (one-time cost per folder)
+    ensure_folder_exists(imap_conn, folder_name)
+    imap_conn.select(f'"{folder_name}"')
+    server_ids = set(get_message_ids_in_folder(imap_conn).values())
+    built.update(server_ids)
+
+    with msg_ids_lock:
+        msg_ids_by_folder.setdefault(folder_name, built)
+        return msg_ids_by_folder[folder_name]
+
+
+def load_manifest(local_path, filename):
+    """Load a manifest JSON file from a backup directory.
+
+    Args:
+        local_path: Directory containing the manifest file
+        filename: Manifest filename (e.g. "labels_manifest.json")
+
+    Returns:
+        Parsed manifest dict, or empty dict if not found or invalid.
+    """
+    manifest_path = os.path.join(local_path, filename)
+    if os.path.exists(manifest_path):
+        try:
+            with open(manifest_path, encoding="utf-8") as f:
+                manifest = json.load(f)
+                safe_print(f"Loaded {filename} with {len(manifest)} entries.")
+                return manifest
+        except Exception as e:
+            safe_print(f"Warning: Could not load {filename}: {e}")
+    return {}

src/utils/imap_compress.py

@@ -0,0 +1,180 @@
+"""IMAP COMPRESS=DEFLATE support (RFC 4978).
+
+Provides a transparent compression wrapper for stdlib imaplib connections.
+After authentication, call ``enable_compression(conn)`` to negotiate
+COMPRESS=DEFLATE with the server.  All subsequent IMAP traffic is then
+compressed/decompressed automatically via zlib.
+
+The implementation follows imaplib's own STARTTLS pattern: send the
+command, and on OK replace ``self.sock`` with a wrapper.
+"""
+
+from __future__ import annotations
+
+import io
+import zlib
+from collections.abc import Callable
+
+
+class _CompressedRawIO(io.RawIOBase):
+    """Minimal RawIOBase adapter so ``makefile()`` can return a BufferedReader
+    that reads through the decompression layer."""
+
+    def __init__(self, compressed_sock: _CompressedSocket):
+        self._sock = compressed_sock
+
+    def readable(self) -> bool:
+        return True
+
+    def readinto(self, b: bytearray | memoryview) -> int:
+        data = self._sock.recv(len(b))
+        if not data:
+            return 0
+        n = len(data)
+        b[:n] = data
+        return n
+
+
+class _CompressedSocket:
+    """Socket wrapper that adds zlib DEFLATE compression/decompression.
+
+    Wraps a real socket (typically SSL) so that:
+    - ``sendall()`` compresses data before sending
+    - ``recv()`` decompresses data after receiving
+
+    Python 3.13's imaplib reads via ``self.sock.recv()`` directly,
+    so intercepting at the socket level is sufficient.
+    """
+
+    def __init__(self, sock, initial_compressed: bytes = b""):
+        self._sock = sock
+        self._compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        self._decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
+        # Buffer for decompressed data not yet consumed by recv()
+        self._recv_buf = b""
+        # If there was read-ahead data in imaplib's buffer after the
+        # COMPRESS OK response, it is already compressed and must be
+        # fed to the decompressor.
+        if initial_compressed:
+            self._recv_buf = self._decompressor.decompress(initial_compressed)
+
+    # -- outgoing (compress) ---------------------------------------------------
+
+    def sendall(self, data: bytes) -> None:
+        compressed = self._compressor.compress(data)
+        compressed += self._compressor.flush(zlib.Z_SYNC_FLUSH)
+        self._sock.sendall(compressed)
+
+    def send(self, data: bytes) -> int:
+        self.sendall(data)
+        return len(data)
+
+    # -- incoming (decompress) -------------------------------------------------
+
+    def recv(self, bufsize: int) -> bytes:
+        # Return buffered decompressed data first
+        if self._recv_buf:
+            chunk = self._recv_buf[:bufsize]
+            self._recv_buf = self._recv_buf[bufsize:]
+            return chunk
+
+        # Read and decompress until we have at least 1 byte or true EOF.
+        # zlib can produce b"" for partial DEFLATE blocks; returning that
+        # would signal EOF to imaplib and break the connection.
+        decompressed = b""
+        while not decompressed:
+            raw = self._sock.recv(max(bufsize, 16384))
+            if not raw:
+                return b""  # real EOF
+            decompressed = self._decompressor.decompress(raw)
+
+        if len(decompressed) > bufsize:
+            self._recv_buf = decompressed[bufsize:]
+            return decompressed[:bufsize]
+        return decompressed
+
+    # -- lifecycle / proxy -----------------------------------------------------
+
+    def shutdown(self, how: int) -> None:
+        self._sock.shutdown(how)
+
+    def close(self) -> None:
+        self._sock.close()
+
+    def makefile(self, *args, **kwargs):
+        # Older Python imaplib reads from conn._file (a file object)
+        # rather than conn.sock.recv() directly.  The file must read
+        # through the decompression layer, not from the raw socket.
+        return io.BufferedReader(_CompressedRawIO(self))
+
+    def __getattr__(self, name: str):
+        return getattr(self._sock, name)
+
+
+def enable_compression(
+    conn,
+    *,
+    log_fn: Callable[[str], None] | None = None,
+) -> bool:
+    """Negotiate COMPRESS=DEFLATE on a raw imaplib IMAP4/IMAP4_SSL connection.
+
+    Must be called **after** authentication (login/XOAUTH2).
+
+    Returns True if compression was successfully enabled, False otherwise
+    (e.g. server does not advertise the capability, or the command failed).
+    """
+    raw_caps = getattr(conn, "capabilities", ())
+    caps = {c.decode() if isinstance(c, bytes) else c for c in raw_caps}
+    if "COMPRESS=DEFLATE" not in caps:
+        if log_fn is not None:
+            log_fn("Deflate compression not supported by server")
+        return False
+
+    try:
+        typ, _data = conn._simple_command("COMPRESS", "DEFLATE")
+    except Exception:
+        if log_fn is not None:
+            log_fn("Deflate compression negotiation failed")
+        return False
+
+    if typ != "OK":
+        if log_fn is not None:
+            log_fn("Deflate compression rejected by server")
+        return False
+
+    # Drain any read-ahead data from imaplib's internal buffers.
+    # After the server sends OK, all subsequent data is compressed.
+    # If recv() read past the OK response, the leftover bytes are
+    # already compressed and must be fed to the decompressor.
+    leftover = b""
+
+    # conn._file (BufferedReader) may have read-ahead bytes beyond
+    # the OK response that are already compressed.  Drain them.
+    old_file = getattr(conn, "_file", None)
+    if old_file is not None:
+        try:
+            buffered = old_file.peek(65536)
+            if isinstance(buffered, (bytes, bytearray)) and buffered:
+                leftover += bytes(buffered)
+        except Exception:
+            pass
+
+    readbuf = getattr(conn, "_readbuf", None)
+    if readbuf:
+        if isinstance(readbuf, (bytes, bytearray)):
+            leftover += bytes(readbuf)
+        else:
+            leftover += b"".join(readbuf)
+        if isinstance(readbuf, bytes):
+            conn._readbuf = b""
+        else:
+            readbuf.clear()
+
+    # Replace the socket (follows the STARTTLS pattern in imaplib).
+    conn.sock = _CompressedSocket(conn.sock, initial_compressed=leftover)
+    conn._file = conn.sock.makefile("rb")
+
+    if log_fn is not None:
+        log_fn("Deflate compression enabled")
+
+    return True

src/utils/imap_retry.py

@@ -0,0 +1,80 @@
+"""
+IMAP Retry Logic
+
+Transparent retry wrapper for IMAP connections that handles transient
+server errors (e.g. Microsoft 365 "Server Busy") with exponential backoff.
+"""
+
+from __future__ import annotations
+
+import time
+
+
+class ConnectionProxy:
+    """Transparent proxy that retries IMAP commands on transient server errors.
+
+    Wraps an imaplib.IMAP4 or IMAP4_SSL connection. For methods in
+    RETRYABLE_METHODS that return (typ, data) tuples, retries on transient
+    errors with exponential backoff.
+    """
+
+    TRANSIENT_PATTERNS = [b"UNAVAILABLE", b"Server Busy", b"try again", b"THROTTLED"]
+
+    # Methods that are safe to retry and return (typ, data)
+    RETRYABLE_METHODS = frozenset(
+        {
+            "uid",
+            "select",
+            "search",
+            "fetch",
+            "append",
+            "store",
+            "list",
+            "create",
+            "expunge",
+            "noop",
+        }
+    )
+
+    def __init__(self, conn, max_retries=3, initial_wait=5, log_fn=print):
+        if max_retries < 1:
+            raise ValueError(f"max_retries must be >= 1, got {max_retries}")
+        if initial_wait < 0:
+            raise ValueError(f"initial_wait must be >= 0, got {initial_wait}")
+        self._conn = conn
+        self._max_retries = max_retries
+        self._initial_wait = initial_wait
+        self._log_fn = log_fn
+
+    @classmethod
+    def _is_transient_error(cls, data):
+        """Check if IMAP response data contains transient error patterns."""
+        for item in data:
+            if isinstance(item, bytes):
+                for pattern in cls.TRANSIENT_PATTERNS:
+                    if pattern in item:
+                        return True
+        return False
+
+    def __getattr__(self, name):
+        attr = getattr(self._conn, name)
+        if name not in self.RETRYABLE_METHODS or not callable(attr):
+            return attr
+
+        def wrapper(*args, **kwargs):
+            last_result = None
+            for attempt in range(self._max_retries):
+                result = attr(*args, **kwargs)
+                if not isinstance(result, tuple) or len(result) < 2:
+                    return result
+                typ, data = result[0], result[1]
+                if typ == "OK" or not self._is_transient_error(data):
+                    return result
+                last_result = result
+                if attempt + 1 < self._max_retries:
+                    wait = self._initial_wait * (2**attempt)  # 5s, 10s, 20s
+                    self._log_fn(f"Server busy, retrying in {wait}s... (attempt {attempt + 1}/{self._max_retries})")
+                    time.sleep(wait)
+            return last_result
+
+        return wrapper

src/utils/restore_cache.py

@@ -67,7 +67,7 @@ def load_dest_index_cache(cache_path: str) -> dict:
         return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
 
 
-def save_dest_index_cache(cache_path: str, cache_data: dict) -> None:
+def save_dest_index_cache(cache_path: str, cache_data: dict) -> bool:
     try:
         tmp_path = f"{cache_path}.tmp"
         with open(tmp_path, "w", encoding="utf-8") as f:
@@ -252,3 +252,62 @@ def record_progress(
             progress_cache_lock,
             log_fn=log_fn,
         )
+
+
+def get_source_data(
+    cache_data: dict,
+    folder_name: str,
+) -> dict:
+    """Retrieve source tracking data for a folder."""
+    folders = cache_data.get("folders", {})
+    if not isinstance(folders, dict):
+        return {}
+    folder_entry = folders.get(folder_name, {})
+    if not isinstance(folder_entry, dict):
+        return {}
+    return folder_entry.get("source_state", {})
+
+
+def set_source_data(
+    cache_data: dict,
+    folder_name: str,
+    uid_validity: int,
+    last_uid: int,
+) -> None:
+    """Update source tracking data for a folder."""
+    if "folders" not in cache_data or not isinstance(cache_data["folders"], dict):
+        cache_data["folders"] = {}
+
+    if folder_name not in cache_data["folders"] or not isinstance(cache_data["folders"][folder_name], dict):
+        cache_data["folders"][folder_name] = {}
+
+    # Ensure folder entry is a dict
+    if not isinstance(cache_data["folders"][folder_name], dict):
+        cache_data["folders"][folder_name] = {}
+
+    cache_data["folders"][folder_name]["source_state"] = {"uid_validity": uid_validity, "last_uid": last_uid}
+
+
+def record_source_progress(
+    *,
+    folder_name: str,
+    uid_validity: int,
+    last_uid: int,
+    progress_cache_path: str | None,
+    progress_cache_data: dict | None,
+    progress_cache_lock: threading.Lock | None,
+    log_fn: Callable[[str], None] | None = None,
+) -> None:
+    """Record source progress (last processed UID) to the cache."""
+    if progress_cache_path and progress_cache_data is not None and progress_cache_lock is not None:
+        with progress_cache_lock:
+            set_source_data(progress_cache_data, folder_name, uid_validity, last_uid)
+
+        # Force save to ensure watermark is persistent
+        maybe_save_dest_index_cache(
+            progress_cache_path,
+            progress_cache_data,
+            progress_cache_lock,
+            force=True,
+            log_fn=log_fn,
+        )

test/auth/test_imap_oauth2.py

@@ -0,0 +1,416 @@
+"""
+Tests for imap_oauth2.py
+
+Tests cover:
+- OAuth2 provider detection
+- Provider dispatch
+- Thread-safe token refresh
+- Token expiration and auth error detection
+"""
+
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+
+from auth import imap_oauth2, oauth2_google, oauth2_microsoft
+
+
+@pytest.fixture(autouse=True)
+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:
+    """Tests for detect_oauth2_provider function."""
+
+    def test_microsoft_outlook(self):
+        """Test detects Microsoft from outlook host."""
+        assert imap_oauth2.detect_oauth2_provider("outlook.office365.com") == "microsoft"
+
+    def test_microsoft_office365(self):
+        """Test detects Microsoft from office365 host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.office365.com") == "microsoft"
+
+    def test_microsoft_mixed_case(self):
+        """Test detects Microsoft case-insensitively."""
+        assert imap_oauth2.detect_oauth2_provider("Outlook.Office365.COM") == "microsoft"
+
+    def test_google_gmail(self):
+        """Test detects Google from gmail host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.gmail.com") == "google"
+
+    def test_google_googlemail(self):
+        """Test detects Google from google host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.google.com") == "google"
+
+    def test_unknown_provider(self):
+        """Test returns None for unrecognized host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.example.com") is None
+
+    def test_unknown_yahoo(self):
+        """Test returns None for Yahoo host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.mail.yahoo.com") is None
+
+
+class TestAcquireOauth2TokenForProvider:
+    """Tests for acquire_oauth2_token_for_provider dispatch function."""
+
+    def test_dispatch_to_microsoft(self):
+        """Test dispatches to Microsoft when provider is 'microsoft'."""
+        with patch.object(oauth2_microsoft, "acquire_token", return_value="ms_token") as mock_ms:
+            result = imap_oauth2.acquire_oauth2_token_for_provider("microsoft", "cid", "user@test.com")
+
+        assert result == "ms_token"
+        mock_ms.assert_called_once_with("cid", "user@test.com")
+
+    def test_dispatch_to_google(self):
+        """Test dispatches to Google when provider is 'google'."""
+        with patch.object(oauth2_google, "acquire_token", return_value="g_token") as mock_g:
+            result = imap_oauth2.acquire_oauth2_token_for_provider("google", "cid", "user@gmail.com", "secret")
+
+        assert result == "g_token"
+        mock_g.assert_called_once_with("cid", "secret")
+
+    def test_google_requires_client_secret(self, capsys):
+        """Test returns None when Google is selected without client_secret."""
+        result = imap_oauth2.acquire_oauth2_token_for_provider("google", "cid", "user@gmail.com")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "--oauth2-client-secret" in captured.out
+
+    def test_unknown_provider(self, capsys):
+        """Test returns None for unknown provider."""
+        result = imap_oauth2.acquire_oauth2_token_for_provider("yahoo", "cid", "user@yahoo.com")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "Unknown OAuth2 provider" in captured.out
+
+
+class TestRefreshOauth2Token:
+    """Tests for thread-safe refresh_oauth2_token function."""
+
+    def test_refreshes_token_and_updates_conf(self):
+        """Test that a new token is acquired and conf["oauth2_token"] is updated."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@test.com",
+                "client_secret": None,
+            },
+        }
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="new_token") as mock_acquire:
+            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result == "new_token"
+        assert conf["oauth2_token"] == "new_token"
+        mock_acquire.assert_called_once_with("microsoft", "client-id", "user@test.com", None)
+
+    def test_skips_refresh_when_token_already_updated(self):
+        """Test that refresh is skipped if another thread already updated the token."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "already_refreshed_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@test.com",
+                "client_secret": None,
+            },
+        }
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider") as mock_acquire:
+            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result == "already_refreshed_token"
+        assert conf["oauth2_token"] == "already_refreshed_token"
+        mock_acquire.assert_not_called()
+
+    def test_returns_none_on_refresh_failure(self):
+        """Test returns None and leaves conf unchanged when refresh fails."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": {
+                "provider": "google",
+                "client_id": "client-id",
+                "email": "user@gmail.com",
+                "client_secret": "secret",
+            },
+        }
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
+            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result is None
+        assert conf["oauth2_token"] == "old_token"
+
+    def test_returns_none_when_no_oauth2_context(self):
+        """Test returns None when oauth2 context is not set."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": None,
+        }
+
+        result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result is None
+        assert conf["oauth2_token"] == "old_token"
+
+    def test_concurrent_threads_only_one_refreshes(self):
+        """Test that only one thread performs the refresh when multiple threads compete."""
+        import threading
+        import time
+
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "expired_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@test.com",
+                "client_secret": None,
+            },
+        }
+        call_count = {"value": 0}
+        barrier = threading.Barrier(3)  # 3 threads
+
+        def slow_acquire(provider, client_id, email, client_secret):
+            call_count["value"] += 1
+            time.sleep(0.05)  # Simulate network delay
+            return "fresh_token"
+
+        def thread_func():
+            barrier.wait()  # Ensure all threads start at the same time
+            imap_oauth2.refresh_oauth2_token(conf, "expired_token")
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", side_effect=slow_acquire):
+            threads = [threading.Thread(target=thread_func) for _ in range(3)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+
+        # Only one thread should have called acquire (the first to get the lock).
+        # 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
+
+
+class TestAcquireToken:
+    """Tests for the acquire_token convenience function."""
+
+    def test_returns_token_and_provider_for_microsoft(self, capsys):
+        """Test successful token acquisition for Microsoft host."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="ms_token"):
+            token, provider = imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com")
+
+        assert token == "ms_token"
+        assert provider == "microsoft"
+        out = capsys.readouterr().out
+        assert "Acquiring OAuth2 token (microsoft)" in out
+        assert "OAuth2 token acquired successfully" in out
+
+    def test_returns_token_and_provider_for_google(self, capsys):
+        """Test successful token acquisition for Google host."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="g_token"):
+            token, provider = imap_oauth2.acquire_token(
+                "imap.gmail.com", "cid", "user@gmail.com", client_secret="secret"
+            )
+
+        assert token == "g_token"
+        assert provider == "google"
+
+    def test_label_appears_in_messages(self, capsys):
+        """Test that label is included in status messages."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="tok"):
+            imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com", label="source")
+
+        out = capsys.readouterr().out
+        assert "Acquiring OAuth2 token for source (microsoft)" in out
+        assert "Source OAuth2 token acquired successfully" in out
+
+    def test_exits_on_unrecognized_host(self):
+        """Test sys.exit(1) when provider cannot be detected from host."""
+        with pytest.raises(SystemExit) as exc_info:
+            imap_oauth2.acquire_token("imap.example.com", "cid", "user@example.com")
+        assert exc_info.value.code == 1
+
+    def test_exits_on_token_acquisition_failure(self):
+        """Test sys.exit(1) when token acquisition returns None."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
+            with pytest.raises(SystemExit) as exc_info:
+                imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com")
+        assert exc_info.value.code == 1
+
+    def test_exit_message_includes_label(self, capsys):
+        """Test that failure message includes the label when provided."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
+            with pytest.raises(SystemExit):
+                imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com", label="destination")
+
+        out = capsys.readouterr().out
+        assert "Failed to acquire OAuth2 token for destination" in out
+
+    def test_exit_message_without_label(self, capsys):
+        """Test that failure message works without a label."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
+            with pytest.raises(SystemExit):
+                imap_oauth2.acquire_token("imap.gmail.com", "cid", "user@gmail.com", client_secret="s")
+
+        out = capsys.readouterr().out
+        assert "Error: Failed to acquire OAuth2 token." in out
+
+    def test_passes_client_secret_to_provider(self):
+        """Test that client_secret is forwarded to acquire_oauth2_token_for_provider."""
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="tok") as mock_acq:
+            imap_oauth2.acquire_token("imap.gmail.com", "cid", "user@gmail.com", client_secret="my_secret")
+
+        mock_acq.assert_called_once_with("google", "cid", "user@gmail.com", "my_secret")
+
+
+class TestAuthDescription:
+    """Tests for the auth_description function."""
+
+    def test_microsoft_provider(self):
+        """Test returns OAuth2 description for Microsoft provider."""
+        assert imap_oauth2.auth_description("microsoft") == "OAuth2/microsoft (XOAUTH2)"
+
+    def test_google_provider(self):
+        """Test returns OAuth2 description for Google provider."""
+        assert imap_oauth2.auth_description("google") == "OAuth2/google (XOAUTH2)"
+
+    def test_none_provider(self):
+        """Test returns Basic description when provider is None."""
+        assert imap_oauth2.auth_description(None) == "Basic (password)"
+
+    def test_falsy_provider(self):
+        """Test returns Basic description for any falsy value."""
+        assert imap_oauth2.auth_description("") == "Basic (password)"
+        assert imap_oauth2.auth_description(0) == "Basic (password)"
+        assert imap_oauth2.auth_description(False) == "Basic (password)"

test/auth/test_oauth2_google.py

@@ -0,0 +1,169 @@
+"""
+Tests for oauth2_google.py
+
+Tests cover:
+- Token acquisition using installed app flow
+- Credentials caching and silent token refresh
+"""
+
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+
+from auth import oauth2_google
+
+
+@pytest.fixture(autouse=True)
+def clear_caches():
+    """Clear module-level caches between tests."""
+    oauth2_google._creds_cache.clear()
+    yield
+    oauth2_google._creds_cache.clear()
+
+
+class TestAcquireToken:
+    """Tests for acquire_token function."""
+
+    def test_successful_token(self):
+        """Test successful Google token acquisition."""
+        mock_credentials = MagicMock()
+        mock_credentials.token = "google_test_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = oauth2_google.acquire_token("client-id", "client-secret")
+
+        assert result == "google_test_token"
+
+    def test_missing_library(self):
+        """Test exits when google-auth-oauthlib is not installed."""
+        with patch.dict("sys.modules", {"google_auth_oauthlib": None, "google_auth_oauthlib.flow": None}):
+            with pytest.raises(SystemExit):
+                oauth2_google.acquire_token("client-id", "client-secret")
+
+    def test_no_token_returned(self):
+        """Test returns None when credentials have no token."""
+        mock_credentials = MagicMock()
+        mock_credentials.token = None
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = oauth2_google.acquire_token("client-id", "client-secret")
+
+        assert result is None
+
+    def test_credentials_cached_on_first_call(self):
+        """Test Google credentials are cached after first call."""
+        mock_credentials = MagicMock()
+        mock_credentials.token = "google_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            oauth2_google.acquire_token("client-id", "client-secret")
+
+        assert ("client-id", "client-secret") in oauth2_google._creds_cache
+
+    def test_cached_credentials_refreshed_on_second_call(self):
+        """Test second call refreshes cached credentials without opening browser."""
+        # Pre-populate cache with credentials that have a refresh token
+        mock_creds = MagicMock()
+        mock_creds.refresh_token = "refresh_tok"
+        mock_creds.token = "refreshed_google_token"
+        oauth2_google._creds_cache[("client-id", "client-secret")] = mock_creds
+
+        mock_request_module = MagicMock()
+        with patch.dict(
+            "sys.modules",
+            {
+                "google": MagicMock(),
+                "google.auth": MagicMock(),
+                "google.auth.transport": MagicMock(),
+                "google.auth.transport.requests": mock_request_module,
+            },
+        ):
+            result = oauth2_google.acquire_token("client-id", "client-secret")
+
+        assert result == "refreshed_google_token"
+        # Verify refresh was called
+        mock_creds.refresh.assert_called_once()
+
+    def test_falls_back_to_browser_if_refresh_fails(self):
+        """Test falls back to full auth flow if cached token refresh fails."""
+        # Pre-populate cache with credentials whose refresh fails
+        mock_creds = MagicMock()
+        mock_creds.refresh_token = "refresh_tok"
+        mock_creds.refresh.side_effect = Exception("Refresh failed")
+        oauth2_google._creds_cache[("client-id", "client-secret")] = mock_creds
+
+        # Set up the full auth flow
+        mock_credentials = MagicMock()
+        mock_credentials.token = "new_browser_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = oauth2_google.acquire_token("client-id", "client-secret")
+
+        assert result == "new_browser_token"
+
+    def test_no_refresh_if_no_refresh_token(self):
+        """Test falls back to browser if cached credentials have no refresh token."""
+        # Pre-populate cache with credentials that have no refresh token
+        mock_creds = MagicMock()
+        mock_creds.refresh_token = None
+        oauth2_google._creds_cache[("client-id", "client-secret")] = mock_creds
+
+        # Set up the full auth flow
+        mock_credentials = MagicMock()
+        mock_credentials.token = "new_browser_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = oauth2_google.acquire_token("client-id", "client-secret")
+
+        assert result == "new_browser_token"
+        # Refresh should not have been called
+        mock_creds.refresh.assert_not_called()

test/auth/test_oauth2_microsoft.py

@@ -0,0 +1,290 @@
+"""
+Tests for oauth2_microsoft.py
+
+Tests cover:
+- Tenant discovery from email domain
+- Token acquisition using MSAL device code flow
+- MSAL app caching and silent token refresh
+"""
+
+import json
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+
+from auth import oauth2_microsoft
+from conftest import temp_env
+from mock_oauth_server import MOCK_TENANT_ID
+
+
+@pytest.fixture(autouse=True)
+def clear_caches():
+    """Clear module-level caches between tests."""
+    oauth2_microsoft._msal_app_cache.clear()
+    oauth2_microsoft._tenant_cache.clear()
+    yield
+    oauth2_microsoft._msal_app_cache.clear()
+    oauth2_microsoft._tenant_cache.clear()
+
+
+class TestDiscoverTenant:
+    """Tests for discover_tenant function."""
+
+    def test_successful_discovery(self):
+        """Test successful tenant ID extraction from OpenID config."""
+        tenant_id = "12345678-abcd-ef01-2345-67890abcdef0"
+        data = {
+            "issuer": f"https://sts.windows.net/{tenant_id}/",
+            "authorization_endpoint": f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/authorize",
+        }
+
+        with patch.object(oauth2_microsoft, "_fetch_json_https", return_value=data):
+            result = oauth2_microsoft.discover_tenant("user@contoso.com")
+
+        assert result == tenant_id
+
+    def test_domain_extraction(self):
+        """Test that domain is correctly extracted from email."""
+        data = {"issuer": "https://sts.windows.net/abcdef01-2345-6789-abcd-ef0123456789/"}
+
+        with patch.object(oauth2_microsoft, "_fetch_json_https", return_value=data) as mock_fetch:
+            oauth2_microsoft.discover_tenant("user@example.org")
+            host, path = mock_fetch.call_args[0][0], mock_fetch.call_args[0][1]
+            assert host == "login.microsoftonline.com"
+            assert "example.org" in path
+
+    def test_network_error(self, capsys):
+        """Test returns None on network error."""
+        with patch.object(oauth2_microsoft, "_fetch_json_https", side_effect=OSError("Connection refused")):
+            result = oauth2_microsoft.discover_tenant("user@invalid.example")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "Could not discover" in captured.out
+
+    def test_invalid_json(self, capsys):
+        """Test returns None on invalid JSON response."""
+        with patch.object(
+            oauth2_microsoft,
+            "_fetch_json_https",
+            side_effect=json.JSONDecodeError("Expecting value", "not json", 0),
+        ):
+            result = oauth2_microsoft.discover_tenant("user@test.com")
+
+        assert result is None
+
+    def test_no_tenant_in_issuer(self, capsys):
+        """Test returns None when issuer has no tenant GUID."""
+        data = {"issuer": "https://sts.windows.net/not-a-guid/"}
+
+        with patch.object(oauth2_microsoft, "_fetch_json_https", return_value=data):
+            result = oauth2_microsoft.discover_tenant("user@test.com")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "Could not extract tenant ID" in captured.out
+
+    def test_tenant_caching(self):
+        """Test that discovered tenant IDs are cached."""
+        tenant_id = "12345678-abcd-ef01-2345-67890abcdef0"
+        data = {"issuer": f"https://sts.windows.net/{tenant_id}/"}
+
+        with patch.object(oauth2_microsoft, "_fetch_json_https", return_value=data) as mock_fetch:
+            # First call
+            result1 = oauth2_microsoft.discover_tenant("user@example.com")
+            # Second call with same domain
+            result2 = oauth2_microsoft.discover_tenant("another@example.com")
+
+        assert result1 == tenant_id
+        assert result2 == tenant_id
+        # Should only fetch once due to caching
+        assert mock_fetch.call_count == 1
+
+    def test_discovery_env_override(self, mock_oauth_server):
+        """Test discovery uses OAUTH2_MICROSOFT_DISCOVERY_URL when set."""
+        with temp_env({"OAUTH2_MICROSOFT_DISCOVERY_URL": mock_oauth_server}):
+            result = oauth2_microsoft.discover_tenant("user@example.org")
+
+        assert result == MOCK_TENANT_ID
+
+
+class TestAcquireToken:
+    """Tests for acquire_token function."""
+
+    def test_successful_token(self):
+        """Test successful token acquisition with auto-discovery."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC123", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "test_token"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                result = oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+            assert result == "test_token"
+
+    def test_custom_authority_url(self):
+        """Test custom authority URL from environment variable."""
+        custom_base = "https://custom.login.example.com"
+        tenant_id = "tenant-123"
+        expected_authority = f"{custom_base}/{tenant_id}"
+
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value=tenant_id):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "msg"}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                with temp_env({"OAUTH2_MICROSOFT_AUTHORITY_BASE_URL": custom_base}):
+                    oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+            mock_msal.PublicClientApplication.assert_called_with("client-id", authority=expected_authority)
+
+    def test_tenant_discovery_failure(self, capsys):
+        """Test returns None when tenant discovery fails."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value=None):
+            result = oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+        assert result is None
+
+    def test_cached_token(self):
+        """Test returns cached token when available."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_account = {"username": "user@test.com"}
+            mock_app.get_accounts.return_value = [mock_account]
+            mock_app.acquire_token_silent.return_value = {"access_token": "cached_token"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                result = oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+            assert result == "cached_token"
+
+    def test_msal_app_cached_on_first_call(self):
+        """Test MSAL app is cached after first call."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token1"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+            assert ("client-id", "tenant-123") in oauth2_microsoft._msal_app_cache
+
+    def test_cached_app_reused_on_second_call(self):
+        """Test second call reuses cached MSAL app instead of creating new one."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token1"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+                # Second call — simulate cached token available (refresh token worked)
+                mock_account = {"username": "user@test.com"}
+                mock_app.get_accounts.return_value = [mock_account]
+                mock_app.acquire_token_silent.return_value = {"access_token": "refreshed_token"}
+
+                result = oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+            assert result == "refreshed_token"
+            # PublicClientApplication should only have been called once (first call)
+            assert mock_msal.PublicClientApplication.call_count == 1
+
+    def test_missing_msal_library(self):
+        """Test exits when msal package is not installed."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value="tenant-123"):
+            with patch.dict("sys.modules", {"msal": None}):
+                with pytest.raises(SystemExit):
+                    oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+    def test_no_token_in_response(self):
+        """Test returns None when MSAL response has no access_token."""
+        with patch.object(oauth2_microsoft, "discover_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {}  # No access_token
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                result = oauth2_microsoft.acquire_token("client-id", "user@test.com")
+
+            assert result is None
+
+
+class TestFetchJsonHttps:
+    """Tests for internal _fetch_json_https function."""
+
+    def test_invalid_host_raises_value_error(self):
+        """Test invalid host raises ValueError."""
+        # Test empty host
+        with pytest.raises(ValueError, match="Invalid host"):
+            oauth2_microsoft._fetch_json_https("", "/path")
+
+        # Test None host
+        with pytest.raises(ValueError, match="Invalid host"):
+            oauth2_microsoft._fetch_json_https(None, "/path")
+
+        # Test host with newline
+        with pytest.raises(ValueError, match="Invalid host"):
+            oauth2_microsoft._fetch_json_https("api.example.com\n", "/path")
+
+    def test_path_normalization(self):
+        """Test path missing leading slash is corrected."""
+        mock_response = MagicMock()
+        mock_response.status = 200
+        mock_response.read.return_value = b'{"key": "value"}'
+
+        mock_conn = MagicMock()
+        mock_conn.getresponse.return_value = mock_response
+
+        with patch("http.client.HTTPSConnection", return_value=mock_conn):
+            # Pass path without '/'
+            result = oauth2_microsoft._fetch_json_https("api.example.com", "my/resource")
+
+            # Check request called with normalized path
+            mock_conn.request.assert_called_with("GET", "/my/resource", headers={"Accept": "application/json"})
+            assert result == {"key": "value"}
+
+    def test_https_ssl_context(self):
+        """Test SSL context is created and passed to HTTPSConnection."""
+        mock_response = MagicMock()
+        mock_response.status = 200
+        mock_response.read.return_value = b"{}"
+
+        mock_conn = MagicMock()
+        mock_conn.getresponse.return_value = mock_response
+
+        mock_context = MagicMock()
+
+        with patch("ssl.create_default_context", return_value=mock_context) as mock_create_context:
+            with patch("http.client.HTTPSConnection", return_value=mock_conn) as mock_https:
+                oauth2_microsoft._fetch_json_https("example.com", "/")
+
+                # Verify SSL context creation
+                mock_create_context.assert_called_once()
+
+                # Verify context passed to connection
+                mock_https.assert_called_with("example.com", timeout=10, context=mock_context)

test/conftest.py

@@ -7,6 +7,7 @@ import os
 import socket
 import sys
 import time
+from contextlib import contextmanager
 
 import pytest
 
@@ -15,6 +16,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../s
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tools")))
 
 from mock_imap_server import start_server_thread
+from mock_oauth_server import start_server_thread as start_oauth_server_thread
 
 
 def get_free_port():
@@ -31,11 +33,19 @@ def create_server_pair(src_data=None, dest_data=None):
     while p2 == p1:
         p2 = get_free_port()
 
-    src_t, src_s = start_server_thread(p1, src_data)
-    dest_t, dest_s = start_server_thread(p2, dest_data)
+    src_s, p1_actual = start_server_thread(p1, src_data)
+    dest_s, p2_actual = start_server_thread(p2, dest_data)
     time.sleep(0.3)
 
-    return (src_t, src_s, p1), (dest_t, dest_s, p2)
+    # We don't have direct access to thread object anymore as it is managed internally by TCPServer/ThreadingMixin or start_server_thread wrapper?
+    # Actually start_server_thread in tools/mock_imap_server.py now returns (server, port).
+    # The server object is a ThreadingMixIn so it handles per-request threads, but the main serve_forever loop is in a thread we created.
+    # But wait, my modified start_server_thread does: t.start(), return server, actual_port.
+    # It dropped returning 't'. I need to fix that or update consumers to not need 't'.
+    # Shutdown() stops serve_forever loop. join() is nice but maybe not strictly required if we trust shutdown.
+    # However, to avoid ResourceWarnings, we might want 't'.
+
+    return (None, src_s, p1_actual), (None, dest_s, p2_actual)
 
 
 def shutdown_server_pair(src_tuple, dest_tuple):
@@ -43,9 +53,13 @@ def shutdown_server_pair(src_tuple, dest_tuple):
     src_t, src_s, _ = src_tuple
     dest_t, dest_s, _ = dest_tuple
     src_s.shutdown()
+    src_s.server_close()  # Explicitly close socket
     dest_s.shutdown()
-    src_t.join(timeout=2)
-    dest_t.join(timeout=2)
+    dest_s.server_close()
+    if src_t:
+        src_t.join(timeout=2)
+    if dest_t:
+        dest_t.join(timeout=2)
 
 
 @pytest.fixture
@@ -80,22 +94,47 @@ def single_mock_server():
 
     def _create(initial_data=None):
         port = get_free_port()
-        thread, server = start_server_thread(port, initial_data)
+        server, actual_port = start_server_thread(port, initial_data)
         time.sleep(0.3)
-        servers.append((thread, server))
-        return server, port
+        servers.append(server)
+        return server, actual_port
 
     yield _create
 
-    for thread, server in servers:
+    for server in servers:
         server.shutdown()
-        thread.join(timeout=2)
+        server.server_close()
+
+
+@contextmanager
+def temp_env(env):
+    original = os.environ.copy()
+    os.environ.clear()
+    os.environ.update(env)
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(original)
+
+
+@contextmanager
+def temp_argv(args):
+    original = sys.argv[:]
+    sys.argv = list(args)
+    try:
+        yield
+    finally:
+        sys.argv = original
 
 
 @pytest.fixture(autouse=True)
-def clean_sys_argv(monkeypatch):
+def clean_sys_argv():
     """Ensure sys.argv is clean for all tests."""
-    monkeypatch.setattr(sys, "argv", ["test_script.py"])
+    original = sys.argv[:]
+    sys.argv = ["test_script.py"]
+    yield
+    sys.argv = original
 
 
 def make_mock_connection(src_port, dest_port, src_user="src_user", dest_user="dest_user"):
@@ -111,7 +150,7 @@ def make_mock_connection(src_port, dest_port, src_user="src_user", dest_user="de
         else:
             raise ValueError(f"Unknown user: {user}")
         c = imaplib.IMAP4("localhost", port)
-        c.login(user, pwd)
+        c.login(user, pwd or "")
         return c
 
     return mock_conn
@@ -124,7 +163,31 @@ def make_single_mock_connection(port):
 
     def mock_conn(host, user, pwd, oauth2_token=None):
         c = imaplib.IMAP4("localhost", port)
-        c.login(user, pwd)
+        c.login(user, pwd or "")
         return c
 
     return mock_conn
+
+
+@pytest.fixture
+def mock_oauth_server():
+    """Starts a mock OAuth2 server for token and discovery endpoints."""
+    thread, server = start_oauth_server_thread(0)
+    host, port = server.server_address
+    base_url = f"http://{host}:{port}"
+
+    yield base_url
+
+    server.shutdown()
+    thread.join(timeout=2)
+
+
+__all__ = [
+    "mock_server_factory",
+    "single_mock_server",
+    "make_mock_connection",
+    "make_single_mock_connection",
+    "mock_oauth_server",
+    "temp_env",
+    "temp_argv",
+]

test/core/test_imap_session.py

@@ -0,0 +1,341 @@
+"""
+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")))
+
+from core import imap_session
+
+
+class TestBuildImapConf:
+    """Tests for build_imap_conf function."""
+
+    def test_password_auth_returns_correct_dict(self):
+        """Test builds correct config for password authentication."""
+        conf = imap_session.build_imap_conf("imap.example.com", "user@example.com", "pass123")
+
+        assert conf["host"] == "imap.example.com"
+        assert conf["user"] == "user@example.com"
+        assert conf["password"] == "pass123"
+        assert conf["oauth2_token"] is None
+        assert conf["oauth2"] is None
+
+    def test_oauth2_auth_acquires_token(self):
+        """Test acquires token and builds oauth2 config when client_id is provided."""
+        with patch.object(
+            imap_session.imap_oauth2, "acquire_token", return_value=("my_token", "microsoft")
+        ) as mock_acquire:
+            conf = imap_session.build_imap_conf(
+                "outlook.office365.com",
+                "user@example.com",
+                "pass",
+                client_id="cid",
+                client_secret="csecret",
+                label="source",
+            )
+
+        mock_acquire.assert_called_once_with("outlook.office365.com", "cid", "user@example.com", "csecret", "source")
+        assert conf["host"] == "outlook.office365.com"
+        assert conf["user"] == "user@example.com"
+        assert conf["password"] == "pass"
+        assert conf["oauth2_token"] == "my_token"
+        assert conf["oauth2"] == {
+            "provider": "microsoft",
+            "client_id": "cid",
+            "email": "user@example.com",
+            "client_secret": "csecret",
+        }
+
+    def test_no_client_id_skips_oauth2(self):
+        """Test that oauth2 is None when client_id is not provided."""
+        with patch.object(imap_session.imap_oauth2, "acquire_token") as mock_acquire:
+            conf = imap_session.build_imap_conf("imap.example.com", "user", "pass")
+
+        mock_acquire.assert_not_called()
+        assert conf["oauth2"] is None
+        assert conf["oauth2_token"] is None
+
+    def test_empty_client_id_skips_oauth2(self):
+        """Test that empty string client_id is treated as no OAuth2."""
+        with patch.object(imap_session.imap_oauth2, "acquire_token") as mock_acquire:
+            conf = imap_session.build_imap_conf("imap.example.com", "user", "pass", client_id="")
+
+        mock_acquire.assert_not_called()
+        assert conf["oauth2"] is None
+
+    def test_none_client_secret_passed_through(self):
+        """Test that client_secret=None is correctly passed to acquire_token and stored."""
+        with patch.object(imap_session.imap_oauth2, "acquire_token", return_value=("tok", "microsoft")):
+            conf = imap_session.build_imap_conf("outlook.office365.com", "user@example.com", "pass", client_id="cid")
+
+        assert conf["oauth2"]["client_secret"] is None
+
+    def test_label_forwarded_to_acquire_token(self):
+        """Test that label is passed through to acquire_token."""
+        with patch.object(imap_session.imap_oauth2, "acquire_token", return_value=("tok", "google")) as mock_acquire:
+            imap_session.build_imap_conf(
+                "imap.gmail.com", "user@gmail.com", "pass", client_id="cid", client_secret="sec", label="destination"
+            )
+
+        mock_acquire.assert_called_once_with("imap.gmail.com", "cid", "user@gmail.com", "sec", "destination")
+
+
+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()

test/providers/test_provider_exchange.py

@@ -0,0 +1,96 @@
+"""
+Tests for provider_exchange.py
+
+Tests cover:
+- Exchange special folder detection
+- EXCHANGE_SKIP_FOLDERS constant validation
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+
+from providers import provider_exchange
+
+
+class TestIsSpecialFolder:
+    """Tests for is_special_folder function."""
+
+    def test_suggested_contacts_is_special(self):
+        """Test that Suggested Contacts is detected as special folder."""
+        assert provider_exchange.is_special_folder("Suggested Contacts") is True
+
+    def test_conversation_history_is_special(self):
+        """Test that Conversation History is detected as special folder."""
+        assert provider_exchange.is_special_folder("Conversation History") is True
+
+    def test_calendar_is_special(self):
+        """Test that Calendar is detected as special folder."""
+        assert provider_exchange.is_special_folder("Calendar") is True
+
+    def test_contacts_is_special(self):
+        """Test that Contacts is detected as special folder."""
+        assert provider_exchange.is_special_folder("Contacts") is True
+
+    def test_inbox_is_not_special(self):
+        """Test that INBOX is not a special folder."""
+        assert provider_exchange.is_special_folder("INBOX") is False
+
+    def test_sent_items_is_not_special(self):
+        """Test that Sent Items is not a special folder."""
+        assert provider_exchange.is_special_folder("Sent Items") is False
+
+    def test_drafts_is_not_special(self):
+        """Test that Drafts is not a special folder."""
+        assert provider_exchange.is_special_folder("Drafts") is False
+
+    def test_user_folder_is_not_special(self):
+        """Test that user-created folders are not special."""
+        assert provider_exchange.is_special_folder("Work") is False
+        assert provider_exchange.is_special_folder("Personal") is False
+
+    def test_case_sensitive_matching(self):
+        """Test that matching is case-sensitive."""
+        assert provider_exchange.is_special_folder("suggested contacts") is False
+        assert provider_exchange.is_special_folder("CALENDAR") is False
+
+    def test_empty_string(self):
+        """Test handling of empty string."""
+        assert provider_exchange.is_special_folder("") is False
+
+    def test_none_value(self):
+        """Test handling of None value."""
+        # This test documents current behavior - function expects string
+        # In practice, callers should handle None before calling
+        try:
+            result = provider_exchange.is_special_folder(None)
+            # If it doesn't raise, check the result
+            assert result is False
+        except (TypeError, AttributeError):
+            # Expected if None is not handled
+            pass
+
+
+class TestExchangeConstants:
+    """Tests for Exchange constants."""
+
+    def test_exchange_skip_folders_constant(self):
+        """Test that EXCHANGE_SKIP_FOLDERS contains expected folders."""
+        assert "Suggested Contacts" in provider_exchange.EXCHANGE_SKIP_FOLDERS
+        assert "Conversation History" in provider_exchange.EXCHANGE_SKIP_FOLDERS
+        assert "Calendar" in provider_exchange.EXCHANGE_SKIP_FOLDERS
+        assert "Contacts" in provider_exchange.EXCHANGE_SKIP_FOLDERS
+
+    def test_exchange_skip_folders_count(self):
+        """Test that EXCHANGE_SKIP_FOLDERS has expected number of entries."""
+        # This test will catch if folders are accidentally added/removed
+        assert len(provider_exchange.EXCHANGE_SKIP_FOLDERS) == 4
+
+    def test_inbox_not_in_skip_folders(self):
+        """Test that INBOX is not in EXCHANGE_SKIP_FOLDERS."""
+        assert "INBOX" not in provider_exchange.EXCHANGE_SKIP_FOLDERS
+
+    def test_sent_items_not_in_skip_folders(self):
+        """Test that Sent Items is not in EXCHANGE_SKIP_FOLDERS."""
+        assert "Sent Items" not in provider_exchange.EXCHANGE_SKIP_FOLDERS

test/providers/test_provider_gmail.py

@@ -0,0 +1,286 @@
+"""
+Tests for provider_gmail.py
+
+Tests cover:
+- Gmail label folder detection
+- Folder name to label name conversion
+- Label name to folder path conversion
+- Gmail label index building
+"""
+
+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")))
+
+from providers import provider_gmail
+
+
+class TestIsLabelFolder:
+    """Tests for is_label_folder function."""
+
+    def test_user_labels(self):
+        """Test that user-created labels are correctly identified."""
+        assert provider_gmail.is_label_folder("Work") is True
+        assert provider_gmail.is_label_folder("Personal") is True
+        assert provider_gmail.is_label_folder("Projects/2024") is True
+
+    def test_inbox_is_label(self):
+        """Test that INBOX is considered a label."""
+        assert provider_gmail.is_label_folder("INBOX") is True
+
+    def test_gmail_sent_and_starred_labels(self):
+        """Test that Gmail Sent and Starred are labels worth preserving."""
+        assert provider_gmail.is_label_folder("[Gmail]/Sent Mail") is True
+        assert provider_gmail.is_label_folder("[Gmail]/Starred") is True
+
+    def test_gmail_system_folders_excluded(self):
+        """Test that Gmail system folders are not considered labels."""
+        assert provider_gmail.is_label_folder("[Gmail]/All Mail") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Spam") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Trash") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Drafts") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Bin") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Important") is False
+
+    def test_other_gmail_folders_excluded(self):
+        """Test that other [Gmail]/ folders are not considered labels."""
+        assert provider_gmail.is_label_folder("[Gmail]/Archive") is False
+
+    def test_empty_string(self):
+        """Test handling of empty string."""
+        assert provider_gmail.is_label_folder("") is True  # Empty string doesn't start with [Gmail]/
+
+    def test_nested_labels(self):
+        """Test nested user labels."""
+        assert provider_gmail.is_label_folder("Work/Projects") is True
+        assert provider_gmail.is_label_folder("Personal/Family/Photos") is True
+
+
+class TestFolderToLabel:
+    """Tests for folder_to_label function."""
+
+    def test_inbox_unchanged(self):
+        """Test that INBOX stays as INBOX."""
+        assert provider_gmail.folder_to_label("INBOX") == "INBOX"
+
+    def test_gmail_sent_mail(self):
+        """Test conversion of [Gmail]/Sent Mail."""
+        assert provider_gmail.folder_to_label("[Gmail]/Sent Mail") == "Sent Mail"
+
+    def test_gmail_starred(self):
+        """Test conversion of [Gmail]/Starred."""
+        assert provider_gmail.folder_to_label("[Gmail]/Starred") == "Starred"
+
+    def test_gmail_drafts(self):
+        """Test conversion of [Gmail]/Drafts."""
+        assert provider_gmail.folder_to_label("[Gmail]/Drafts") == "Drafts"
+
+    def test_user_label_unchanged(self):
+        """Test that user labels remain unchanged."""
+        assert provider_gmail.folder_to_label("Work") == "Work"
+        assert provider_gmail.folder_to_label("Personal") == "Personal"
+
+    def test_nested_label_unchanged(self):
+        """Test that nested user labels remain unchanged."""
+        assert provider_gmail.folder_to_label("Work/Projects") == "Work/Projects"
+
+
+class TestLabelToFolder:
+    """Tests for label_to_folder function."""
+
+    def test_inbox_unchanged(self):
+        """Test that INBOX stays as INBOX."""
+        assert provider_gmail.label_to_folder("INBOX") == "INBOX"
+
+    def test_sent_mail_to_gmail_folder(self):
+        """Test conversion of Sent Mail label to [Gmail]/ folder."""
+        assert provider_gmail.label_to_folder("Sent Mail") == "[Gmail]/Sent Mail"
+
+    def test_starred_to_gmail_folder(self):
+        """Test conversion of Starred label to [Gmail]/ folder."""
+        assert provider_gmail.label_to_folder("Starred") == "[Gmail]/Starred"
+
+    def test_drafts_to_gmail_folder(self):
+        """Test conversion of Drafts label to [Gmail]/ folder."""
+        assert provider_gmail.label_to_folder("Drafts") == "[Gmail]/Drafts"
+
+    def test_important_to_gmail_folder(self):
+        """Test conversion of Important label to [Gmail]/ folder."""
+        assert provider_gmail.label_to_folder("Important") == "[Gmail]/Important"
+
+    def test_user_label_unchanged(self):
+        """Test that user labels remain unchanged."""
+        assert provider_gmail.label_to_folder("Work") == "Work"
+        assert provider_gmail.label_to_folder("Personal") == "Personal"
+
+    def test_nested_label_unchanged(self):
+        """Test that nested user labels remain unchanged."""
+        assert provider_gmail.label_to_folder("Work/Projects") == "Work/Projects"
+
+
+class TestResolveTarget:
+    """Tests for resolve_target function."""
+
+    def test_picks_first_valid_label_as_target(self):
+        target, remaining = provider_gmail.resolve_target(["Work", "Personal"])
+        assert target == "Work"
+        assert remaining == ["Personal"]
+
+    def test_skips_system_folders(self):
+        target, remaining = provider_gmail.resolve_target(["[Gmail]/All Mail", "[Gmail]/Spam", "[Gmail]/Trash", "Work"])
+        assert target == "Work"
+        assert remaining == []
+
+    def test_all_system_labels_returns_unlabeled(self):
+        target, remaining = provider_gmail.resolve_target(["[Gmail]/All Mail", "[Gmail]/Spam"])
+        assert target == "Restored/Unlabeled"
+        assert remaining == []
+
+    def test_empty_labels_returns_unlabeled(self):
+        target, remaining = provider_gmail.resolve_target([])
+        assert target == "Restored/Unlabeled"
+        assert remaining == []
+
+    def test_special_labels_mapped_to_gmail_folders(self):
+        target, remaining = provider_gmail.resolve_target(["Sent Mail", "Work"])
+        assert target == "[Gmail]/Sent Mail"
+        assert remaining == ["Work"]
+
+    def test_multiple_remaining_labels(self):
+        target, remaining = provider_gmail.resolve_target(["Work", "Personal", "Finance"])
+        assert target == "Work"
+        assert remaining == ["Personal", "Finance"]
+
+
+class TestBuildGmailLabelIndex:
+    """Tests for build_gmail_label_index function."""
+
+    def test_builds_index_from_label_folders(self):
+        """Test building label index from multiple label folders."""
+        mock_conn = MagicMock()
+
+        # Mock list_selectable_folders to return a mix of folders
+        def mock_list_folders(conn):
+            return ["INBOX", "[Gmail]/All Mail", "Work", "Personal", "[Gmail]/Sent Mail"]
+
+        # Mock get_message_ids_in_folder to return different messages for each folder
+        folder_messages = {
+            "INBOX": {"1": "<msg1@test.com>", "2": "<msg2@test.com>"},
+            "Work": {"3": "<msg1@test.com>", "4": "<msg3@test.com>"},
+            "Personal": {"5": "<msg2@test.com>", "6": "<msg4@test.com>"},
+            "[Gmail]/Sent Mail": {"7": "<msg1@test.com>"},
+        }
+
+        def mock_get_message_ids(conn):
+            # Determine which folder we're in based on the select call
+            selected_folder = mock_conn.select.call_args[0][0].strip('"')
+            return folder_messages.get(selected_folder, {})
+
+        def mock_safe_print(msg):
+            pass
+
+        from utils import imap_common
+
+        with (
+            patch.object(imap_common, "list_selectable_folders", mock_list_folders),
+            patch.object(imap_common, "get_message_ids_in_folder", mock_get_message_ids),
+        ):
+            result = provider_gmail.build_gmail_label_index(mock_conn, mock_safe_print)
+
+        # Verify the index contains the expected mappings
+        assert "<msg1@test.com>" in result
+        assert "INBOX" in result["<msg1@test.com>"]
+        assert "Work" in result["<msg1@test.com>"]
+        assert "Sent Mail" in result["<msg1@test.com>"]
+
+        assert "<msg2@test.com>" in result
+        assert "INBOX" in result["<msg2@test.com>"]
+        assert "Personal" in result["<msg2@test.com>"]
+
+        assert "<msg3@test.com>" in result
+        assert "Work" in result["<msg3@test.com>"]
+
+        assert "<msg4@test.com>" in result
+        assert "Personal" in result["<msg4@test.com>"]
+
+    def test_excludes_system_folders(self):
+        """Test that system folders like [Gmail]/All Mail are excluded from the index."""
+        mock_conn = MagicMock()
+
+        def mock_list_folders(conn):
+            return ["[Gmail]/All Mail", "[Gmail]/Spam", "[Gmail]/Trash", "Work"]
+
+        def mock_get_message_ids(conn):
+            selected_folder = mock_conn.select.call_args[0][0].strip('"')
+            if selected_folder == "Work":
+                return {"1": "<msg1@test.com>"}
+            return {}
+
+        def mock_safe_print(msg):
+            pass
+
+        from utils import imap_common
+
+        with (
+            patch.object(imap_common, "list_selectable_folders", mock_list_folders),
+            patch.object(imap_common, "get_message_ids_in_folder", mock_get_message_ids),
+        ):
+            result = provider_gmail.build_gmail_label_index(mock_conn, mock_safe_print)
+
+        # Should only have the Work label, system folders should be excluded
+        assert "<msg1@test.com>" in result
+        assert "Work" in result["<msg1@test.com>"]
+        assert "All Mail" not in result.get("<msg1@test.com>", set())
+
+    def test_handles_folder_error_gracefully(self):
+        """Test that errors in individual folders don't stop the entire process."""
+        mock_conn = MagicMock()
+
+        def mock_list_folders(conn):
+            return ["INBOX", "Work"]
+
+        def mock_get_message_ids(conn):
+            selected_folder = mock_conn.select.call_args[0][0].strip('"')
+            if selected_folder == "INBOX":
+                return {"1": "<msg1@test.com>"}
+            # Simulate error for Work folder
+            raise Exception("Failed to fetch Work folder")
+
+        def mock_safe_print(msg):
+            pass
+
+        from utils import imap_common
+
+        with (
+            patch.object(imap_common, "list_selectable_folders", mock_list_folders),
+            patch.object(imap_common, "get_message_ids_in_folder", mock_get_message_ids),
+        ):
+            # Should not raise, should continue and build partial index
+            result = provider_gmail.build_gmail_label_index(mock_conn, mock_safe_print)
+
+        # Should have INBOX data despite Work folder error
+        assert "<msg1@test.com>" in result
+        assert "INBOX" in result["<msg1@test.com>"]
+
+
+class TestGmailConstants:
+    """Tests for Gmail constants."""
+
+    def test_gmail_system_folders_constant(self):
+        """Test that GMAIL_SYSTEM_FOLDERS contains expected folders."""
+        assert "[Gmail]/All Mail" in provider_gmail.GMAIL_SYSTEM_FOLDERS
+        assert "[Gmail]/Spam" in provider_gmail.GMAIL_SYSTEM_FOLDERS
+        assert "[Gmail]/Trash" in provider_gmail.GMAIL_SYSTEM_FOLDERS
+        assert "[Gmail]/Drafts" in provider_gmail.GMAIL_SYSTEM_FOLDERS
+        assert "[Gmail]/Bin" in provider_gmail.GMAIL_SYSTEM_FOLDERS
+        assert "[Gmail]/Important" in provider_gmail.GMAIL_SYSTEM_FOLDERS
+
+    def test_gmail_sent_not_in_system_folders(self):
+        """Test that [Gmail]/Sent Mail is not in GMAIL_SYSTEM_FOLDERS."""
+        assert "[Gmail]/Sent Mail" not in provider_gmail.GMAIL_SYSTEM_FOLDERS
+
+    def test_gmail_starred_not_in_system_folders(self):
+        """Test that [Gmail]/Starred is not in GMAIL_SYSTEM_FOLDERS."""
+        assert "[Gmail]/Starred" not in provider_gmail.GMAIL_SYSTEM_FOLDERS

test/test_backup_imap_emails.py

@@ -1,848 +0,0 @@
-"""
-Tests for backup_imap_emails.py
-
-Tests cover:
-- Basic email backup to local files
-- Incremental backup (skip existing)
-- Multiple folder backup
-- Filename sanitization
-- Configuration validation
-"""
-
-import os
-import sys
-from unittest.mock import MagicMock
-
-import pytest
-
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
-
-import backup_imap_emails
-import imap_common
-from conftest import make_single_mock_connection
-
-
-class TestBackupBasic:
-    """Tests for basic backup functionality."""
-
-    def test_single_email_backup(self, single_mock_server, monkeypatch, tmp_path):
-        """Test backing up a single email to local directory."""
-        src_data = {"INBOX": [b"Subject: Test Email\r\nMessage-ID: <1@test>\r\n\r\nBody content"]}
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        # Check that backup was created
-        inbox_path = tmp_path / "INBOX"
-        assert inbox_path.exists()
-
-        eml_files = list(inbox_path.glob("*.eml"))
-        assert len(eml_files) == 1
-
-        # Check content
-        content = eml_files[0].read_bytes()
-        assert b"Test Email" in content or b"Body content" in content
-
-    def test_multiple_emails_backup(self, single_mock_server, monkeypatch, tmp_path):
-        """Test backing up multiple emails."""
-        src_data = {
-            "INBOX": [
-                b"Subject: Email 1\r\nMessage-ID: <1@test>\r\n\r\nBody 1",
-                b"Subject: Email 2\r\nMessage-ID: <2@test>\r\n\r\nBody 2",
-                b"Subject: Email 3\r\nMessage-ID: <3@test>\r\n\r\nBody 3",
-            ]
-        }
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        inbox_path = tmp_path / "INBOX"
-        eml_files = list(inbox_path.glob("*.eml"))
-        assert len(eml_files) == 3
-
-
-class TestIncrementalBackup:
-    """Tests for incremental backup functionality."""
-
-    def test_skip_existing_emails(self, single_mock_server, monkeypatch, tmp_path):
-        """Test that existing emails are skipped during incremental backup."""
-        src_data = {
-            "INBOX": [
-                b"Subject: Email 1\r\nMessage-ID: <1@test>\r\n\r\nBody 1",
-                b"Subject: Email 2\r\nMessage-ID: <2@test>\r\n\r\nBody 2",
-            ]
-        }
-        _, port = single_mock_server(src_data)
-
-        # Create pre-existing file
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-        existing_file = inbox_path / "1_Email_1.eml"
-        existing_file.write_bytes(b"existing content")
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        # Should still have original content (not overwritten)
-        assert existing_file.read_bytes() == b"existing content"
-
-        # But second email should be backed up
-        eml_files = list(inbox_path.glob("*.eml"))
-        assert len(eml_files) == 2
-
-
-class TestMultipleFolderBackup:
-    """Tests for backing up multiple folders."""
-
-    def test_backup_all_folders(self, single_mock_server, monkeypatch, tmp_path):
-        """Test backing up emails from multiple folders."""
-        src_data = {
-            "INBOX": [b"Subject: Inbox\r\nMessage-ID: <1@test>\r\n\r\nC"],
-            "Sent": [b"Subject: Sent\r\nMessage-ID: <2@test>\r\n\r\nC"],
-            "Archive": [b"Subject: Archive\r\nMessage-ID: <3@test>\r\n\r\nC"],
-        }
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        assert (tmp_path / "INBOX").exists()
-        assert (tmp_path / "Sent").exists()
-        assert (tmp_path / "Archive").exists()
-
-    def test_backup_single_folder(self, single_mock_server, monkeypatch, tmp_path):
-        """Test backing up a specific folder only."""
-        src_data = {
-            "INBOX": [b"Subject: Inbox\r\nMessage-ID: <1@test>\r\n\r\nC"],
-            "Sent": [b"Subject: Sent\r\nMessage-ID: <2@test>\r\n\r\nC"],
-        }
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path), "INBOX"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        assert (tmp_path / "INBOX").exists()
-        # Sent should NOT be backed up
-        assert not (tmp_path / "Sent").exists()
-
-
-class TestEmptyFolderHandling:
-    """Tests for empty folder handling."""
-
-    def test_empty_folder(self, single_mock_server, monkeypatch, tmp_path):
-        """Test handling of empty folders."""
-        src_data = {"INBOX": [], "Empty": []}
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        # Should complete without error
-        backup_imap_emails.main()
-
-
-class TestConfigValidation:
-    """Tests for configuration validation."""
-
-    def test_missing_credentials(self, monkeypatch, capsys, tmp_path):
-        """Test that missing credentials cause exit."""
-        env = {}
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
-
-        with pytest.raises(SystemExit) as exc_info:
-            backup_imap_emails.main()
-
-        assert exc_info.value.code == 2
-
-    def test_missing_dest_path(self, monkeypatch, capsys):
-        """Test that missing destination path causes exit."""
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py"])
-
-        with pytest.raises(SystemExit) as exc_info:
-            backup_imap_emails.main()
-
-        assert exc_info.value.code == 2
-
-
-class TestGetExistingUids:
-    """Tests for get_existing_uids function."""
-
-    def test_empty_directory(self, tmp_path):
-        """Test with empty directory."""
-        result = backup_imap_emails.get_existing_uids(str(tmp_path))
-        assert result == set()
-
-    def test_nonexistent_directory(self):
-        """Test with non-existent directory."""
-        result = backup_imap_emails.get_existing_uids("/nonexistent/path")
-        assert result == set()
-
-    def test_existing_files(self, tmp_path):
-        """Test extraction of UIDs from existing files."""
-        (tmp_path / "1_Subject.eml").touch()
-        (tmp_path / "2_Another.eml").touch()
-        (tmp_path / "100_Long_Subject.eml").touch()
-
-        result = backup_imap_emails.get_existing_uids(str(tmp_path))
-        assert result == {"1", "2", "100"}
-
-    def test_ignore_non_matching_files(self, tmp_path):
-        """Test that non-matching files are ignored."""
-        (tmp_path / "not_an_email.txt").touch()
-        (tmp_path / "random.eml").touch()  # No underscore
-        (tmp_path / "1_Subject.eml").touch()
-
-        result = backup_imap_emails.get_existing_uids(str(tmp_path))
-        assert result == {"1"}
-
-    def test_os_error_handling(self, monkeypatch):
-        """Test handling of OS errors during listing."""
-
-        def mock_listdir(path):
-            raise OSError("Access denied")
-
-        monkeypatch.setattr(os, "listdir", mock_listdir)
-
-        uids = backup_imap_emails.get_existing_uids("/some/path")
-        assert len(uids) == 0
-
-
-class TestBackupErrorHandling:
-    """Tests for error handling scenarios in backup."""
-
-    def test_connection_error_in_worker(self, monkeypatch, tmp_path):
-        """Test worker handles connection failure gracefully."""
-        # Mock get_imap_connection to fail
-        monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: None)
-
-        # Should return None/Exit without crashing
-        backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), str(tmp_path))
-
-    def test_select_error_in_worker(self, monkeypatch, tmp_path):
-        """Test worker handles SELECT failure."""
-        mock_conn = MagicMock()
-        mock_conn.select.side_effect = Exception("Select error")
-        monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn)
-
-        # Should log error and return
-        backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), str(tmp_path))
-        mock_conn.select.assert_called()
-
-    def test_fetch_body_error(self, monkeypatch, tmp_path):
-        """Test handling of fetch body failure."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = "OK"
-
-        # Fetch body fails
-        mock_conn.uid.return_value = ("NO", [None])
-
-        monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn)
-
-        # Try processing one UID
-        backup_imap_emails.process_batch([b"1"], "INBOX", ("h", "u", "p"), str(tmp_path))
-
-        # File should not exist
-        assert not list(tmp_path.glob("*.eml"))
-
-    def test_write_error(self, monkeypatch, tmp_path):
-        """Test handling of file write error."""
-        mock_conn = MagicMock()
-        monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn)
-
-        # Mock fetched data
-        mock_conn.uid.return_value = ("OK", [(b"1 (RFC822 {10}", b"Content")])
-
-        # Mock open to fail
-        def mock_open(*args, **kwargs):
-            raise OSError("Disk full")
-
-        monkeypatch.setattr("builtins.open", mock_open)
-
-        backup_imap_emails.process_batch([b"1"], "INBOX", ("h", "u", "p"), str(tmp_path))
-
-    def test_folder_creation_error(self, monkeypatch, tmp_path):
-        """Test handling failure to create local folder."""
-
-        def mock_makedirs(path, exist_ok=False):
-            raise OSError("Permission denied")
-
-        monkeypatch.setattr(os, "makedirs", mock_makedirs)
-
-        mock_conn = MagicMock()
-
-        # backup_folder should return early
-        backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p"))
-        mock_conn.select.assert_not_called()
-
-    def test_select_folder_error(self, monkeypatch, tmp_path):
-        """Test handling of select folder failure in main loop."""
-        mock_conn = MagicMock()
-        mock_conn.select.side_effect = Exception("Select failed")
-        monkeypatch.setattr(os, "makedirs", lambda p, exist_ok: None)
-
-        backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p"))
-        mock_conn.uid.assert_not_called()
-
-    def test_search_error(self, monkeypatch, tmp_path):
-        """Test handling of search failure."""
-        mock_conn = MagicMock()
-        mock_conn.uid.return_value = ("NO", [])
-        monkeypatch.setattr(os, "makedirs", lambda p, exist_ok: None)
-
-        backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p"))
-
-    def test_main_makedirs_error(self, monkeypatch, capsys):
-        """Test failure to create main backup directory."""
-        env = {
-            "SRC_IMAP_HOST": "h",
-            "SRC_IMAP_USERNAME": "u",
-            "SRC_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup.py", "--dest-path", "/protected/path"])
-
-        def mock_makedirs(path):
-            raise OSError("No permission")
-
-        monkeypatch.setattr(os, "makedirs", mock_makedirs)
-        monkeypatch.setattr(os.path, "exists", lambda p: False)
-
-        with pytest.raises(SystemExit):
-            backup_imap_emails.main()
-
-        captured = capsys.readouterr()
-        assert "Error creating backup directory" in captured.out
-
-
-class TestGmailLabelsPreservation:
-    """Tests for Gmail labels manifest functionality."""
-
-    def test_is_gmail_label_folder_user_labels(self):
-        """Test that user labels are correctly identified."""
-        assert backup_imap_emails.is_gmail_label_folder("Work") is True
-        assert backup_imap_emails.is_gmail_label_folder("Personal") is True
-        assert backup_imap_emails.is_gmail_label_folder("Projects/2024") is True
-        assert backup_imap_emails.is_gmail_label_folder("INBOX") is True
-
-    def test_is_gmail_label_folder_gmail_labels(self):
-        """Test that Gmail system labels are correctly identified."""
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Sent Mail") is True
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Starred") is True
-
-    def test_is_gmail_label_folder_system_folders(self):
-        """Test that Gmail system folders are excluded."""
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/All Mail") is False
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Spam") is False
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Trash") is False
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Drafts") is False
-        assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Bin") is False
-
-    def test_load_labels_manifest_nonexistent(self, tmp_path):
-        """Test loading manifest when file doesn't exist."""
-        result = backup_imap_emails.load_labels_manifest(str(tmp_path))
-        assert result == {}
-
-    def test_load_labels_manifest_existing(self, tmp_path):
-        """Test loading existing manifest file."""
-        import json
-
-        manifest_data = {
-            "<msg1@test.com>": ["INBOX", "Work"],
-            "<msg2@test.com>": ["Sent Mail", "Personal"],
-        }
-        manifest_path = tmp_path / "labels_manifest.json"
-        manifest_path.write_text(json.dumps(manifest_data))
-
-        result = backup_imap_emails.load_labels_manifest(str(tmp_path))
-        assert result == manifest_data
-
-    def test_load_labels_manifest_invalid_json(self, tmp_path):
-        """Test loading invalid JSON manifest file."""
-        manifest_path = tmp_path / "labels_manifest.json"
-        manifest_path.write_text("not valid json {{{")
-
-        result = backup_imap_emails.load_labels_manifest(str(tmp_path))
-        assert result == {}
-
-    def test_get_message_ids_in_folder(self, monkeypatch):
-        """Test extraction of message IDs from a folder."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.uid.side_effect = [
-            ("OK", [b"1 2 3"]),  # search result
-            (
-                "OK",
-                [
-                    (b"1 (FLAGS (\\Seen) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg1@test.com>\r\n"),
-                    b")",
-                    (b"2 (FLAGS () BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg2@test.com>\r\n"),
-                    b")",
-                    (
-                        b"3 (FLAGS (\\Seen \\Answered) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}",
-                        b"Message-ID: <msg3@test.com>\r\n",
-                    ),
-                    b")",
-                ],
-            ),  # fetch result
-        ]
-
-        result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", None)
-
-        assert "<msg1@test.com>" in result
-        assert "<msg2@test.com>" in result
-        assert "<msg3@test.com>" in result
-
-    def test_get_message_info_in_folder_read_status(self, monkeypatch):
-        """Test extraction of message IDs with read/unread status."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.uid.side_effect = [
-            ("OK", [b"1 2 3"]),  # search result
-            (
-                "OK",
-                [
-                    (b"1 (FLAGS (\\Seen) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg1@test.com>\r\n"),
-                    b")",
-                    (b"2 (FLAGS () BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg2@test.com>\r\n"),
-                    b")",
-                    (
-                        b"3 (FLAGS (\\Seen \\Answered) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}",
-                        b"Message-ID: <msg3@test.com>\r\n",
-                    ),
-                    b")",
-                ],
-            ),  # fetch result
-        ]
-
-        result = backup_imap_emails.get_message_info_in_folder(mock_conn, "INBOX", None)
-
-        assert "<msg1@test.com>" in result
-        assert "\\Seen" in result["<msg1@test.com>"]["flags"]  # Has \Seen flag
-        assert "<msg2@test.com>" in result
-        assert result["<msg2@test.com>"]["flags"] == []  # No flags
-        assert "<msg3@test.com>" in result
-        assert "\\Seen" in result["<msg3@test.com>"]["flags"]  # Has \Seen flag
-        assert "\\Answered" in result["<msg3@test.com>"]["flags"]  # Also has \Answered
-
-    def test_get_message_ids_in_folder_with_progress(self, monkeypatch):
-        """Test extraction of message IDs with progress callback."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.uid.side_effect = [
-            ("OK", [b"1 2 3"]),  # search result
-            (
-                "OK",
-                [
-                    (b"1 (FLAGS (\\Seen) BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg1@test.com>\r\n"),
-                    b")",
-                ],
-            ),  # fetch result
-        ]
-
-        progress_calls = []
-
-        def progress_cb(current, total):
-            progress_calls.append((current, total))
-
-        backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", progress_cb)
-
-        # Progress should have been called
-        assert len(progress_calls) > 0
-        # Last call should show completion
-        assert progress_calls[-1][0] == progress_calls[-1][1]
-
-    def test_get_message_ids_in_folder_select_error(self, monkeypatch):
-        """Test handling of folder select error."""
-        mock_conn = MagicMock()
-        mock_conn.select.side_effect = Exception("Select failed")
-
-        result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX")
-        assert result == set()
-
-    def test_get_message_ids_in_folder_empty(self, monkeypatch):
-        """Test extraction from empty folder."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"0"])
-        mock_conn.uid.return_value = ("OK", [b""])
-
-        result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", None)
-        assert result == set()
-
-    def test_build_labels_manifest(self, monkeypatch, tmp_path):
-        """Test building labels manifest from mock folders."""
-        mock_conn = MagicMock()
-
-        # Mock folder list - simulate Gmail structure
-        mock_conn.list.return_value = (
-            "OK",
-            [
-                b'(\\HasNoChildren) "/" "INBOX"',
-                b'(\\HasNoChildren) "/" "Work"',
-                b'(\\HasNoChildren) "/" "[Gmail]/All Mail"',
-                b'(\\HasNoChildren) "/" "[Gmail]/Sent Mail"',
-            ],
-        )
-
-        # Track which folder is selected
-        folder_data = {
-            "INBOX": {"<msg1@test.com>", "<msg2@test.com>"},
-            "Work": {"<msg1@test.com>"},
-            "[Gmail]/Sent Mail": {"<msg2@test.com>"},
-        }
-
-        # Mock info for All Mail (with flags)
-        all_mail_info = {
-            "<msg1@test.com>": {"flags": ["\\Seen", "\\Flagged"]},
-            "<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):
-            if folder == "[Gmail]/All Mail":
-                return all_mail_info
-            return {}
-
-        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)
-
-        result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path))
-
-        # Check manifest structure (new format with labels and flags)
-        assert "<msg1@test.com>" in result
-        assert "<msg2@test.com>" in result
-        assert "INBOX" in result["<msg1@test.com>"]["labels"]
-        assert "Work" in result["<msg1@test.com>"]["labels"]
-        assert "\\Seen" in result["<msg1@test.com>"]["flags"]
-        assert "\\Flagged" in result["<msg1@test.com>"]["flags"]
-        assert "INBOX" in result["<msg2@test.com>"]["labels"]
-        assert "Sent Mail" in result["<msg2@test.com>"]["labels"]
-        assert result["<msg2@test.com>"]["flags"] == []
-
-        # Check file was saved
-        manifest_path = tmp_path / "labels_manifest.json"
-        assert manifest_path.exists()
-
-    def test_build_labels_manifest_list_error(self, monkeypatch, tmp_path):
-        """Test handling of folder list error."""
-        mock_conn = MagicMock()
-        mock_conn.list.return_value = ("NO", [])
-
-        result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path))
-        assert result == {}
-
-    def test_preserve_labels_flag_integration(self, single_mock_server, monkeypatch, tmp_path):
-        """Test --preserve-labels flag creates manifest."""
-        src_data = {
-            "INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-            "Work": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-            "[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-        }
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["backup_imap_emails.py", "--dest-path", str(tmp_path), "--preserve-labels", "[Gmail]/All Mail"],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        # Check manifest was created
-        manifest_path = tmp_path / "labels_manifest.json"
-        assert manifest_path.exists()
-
-    def test_manifest_only_flag(self, single_mock_server, monkeypatch, tmp_path):
-        """Test --manifest-only flag creates manifest without downloading emails."""
-        src_data = {
-            "INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-            "Work": [b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody"],
-            "[Gmail]/All Mail": [
-                b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody",
-                b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody",
-            ],
-        }
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["backup_imap_emails.py", "--dest-path", str(tmp_path), "--manifest-only"],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        # Should exit with code 0 after creating manifest
-        with pytest.raises(SystemExit) as exc_info:
-            backup_imap_emails.main()
-
-        assert exc_info.value.code == 0
-
-        # Check manifest was created
-        manifest_path = tmp_path / "labels_manifest.json"
-        assert manifest_path.exists()
-
-        # Check NO email folders were created (no download happened)
-        # Only the manifest file should exist
-        items = list(tmp_path.iterdir())
-        assert len(items) == 1
-        assert items[0].name == "labels_manifest.json"
-
-    def test_gmail_system_folders_constant(self):
-        """Test that GMAIL_SYSTEM_FOLDERS contains expected entries."""
-        expected = {
-            "[Gmail]/All Mail",
-            "[Gmail]/Spam",
-            "[Gmail]/Trash",
-            "[Gmail]/Drafts",
-            "[Gmail]/Bin",
-            "[Gmail]/Important",
-        }
-        assert imap_common.GMAIL_SYSTEM_FOLDERS == expected
-
-    def test_gmail_mode_flag(self, single_mock_server, monkeypatch, tmp_path):
-        """Test --gmail-mode flag backs up All Mail and creates manifest."""
-        src_data = {
-            "INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-            "Work": [b"Subject: Work Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-            "[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
-        }
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["backup_imap_emails.py", "--dest-path", str(tmp_path), "--gmail-mode"],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        # Check manifest was created
-        manifest_path = tmp_path / "labels_manifest.json"
-        assert manifest_path.exists()
-
-        # Check only [Gmail]/All Mail folder was backed up (not INBOX or Work)
-        all_mail_folder = tmp_path / "[Gmail]" / "All Mail"
-        assert all_mail_folder.exists()
-
-        # Should NOT have created separate INBOX or Work folders
-        inbox_folder = tmp_path / "INBOX"
-        work_folder = tmp_path / "Work"
-        assert not inbox_folder.exists()
-        assert not work_folder.exists()
-
-
-class TestDeleteOrphanLocalFiles:
-    """Tests for delete_orphan_local_files function."""
-
-    def test_delete_orphan_files(self, tmp_path):
-        """Test that local files not on server are deleted."""
-        # Create a local folder with some .eml files
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-
-        # Create files with UIDs 1, 2, 3
-        (inbox_path / "1_Test_Email.eml").write_bytes(b"content1")
-        (inbox_path / "2_Another_Email.eml").write_bytes(b"content2")
-        (inbox_path / "3_Third_Email.eml").write_bytes(b"content3")
-
-        # Server only has UIDs 1 and 3
-        server_uids = {"1", "3"}
-
-        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
-
-        # UID 2 should be deleted
-        assert deleted == 1
-        assert (inbox_path / "1_Test_Email.eml").exists()
-        assert not (inbox_path / "2_Another_Email.eml").exists()
-        assert (inbox_path / "3_Third_Email.eml").exists()
-
-    def test_delete_multiple_orphans(self, tmp_path):
-        """Test deleting multiple orphan files."""
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-
-        # Create files with UIDs 1-5
-        for i in range(1, 6):
-            (inbox_path / f"{i}_Email_{i}.eml").write_bytes(f"content{i}".encode())
-
-        # Server only has UID 3
-        server_uids = {"3"}
-
-        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
-
-        # 4 files should be deleted (UIDs 1, 2, 4, 5)
-        assert deleted == 4
-        assert not (inbox_path / "1_Email_1.eml").exists()
-        assert not (inbox_path / "2_Email_2.eml").exists()
-        assert (inbox_path / "3_Email_3.eml").exists()
-        assert not (inbox_path / "4_Email_4.eml").exists()
-        assert not (inbox_path / "5_Email_5.eml").exists()
-
-    def test_no_orphans(self, tmp_path):
-        """Test when all local files exist on server."""
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-
-        (inbox_path / "1_Email.eml").write_bytes(b"content")
-        (inbox_path / "2_Email.eml").write_bytes(b"content")
-
-        server_uids = {"1", "2"}
-
-        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
-
-        assert deleted == 0
-        assert (inbox_path / "1_Email.eml").exists()
-        assert (inbox_path / "2_Email.eml").exists()
-
-    def test_nonexistent_folder(self, tmp_path):
-        """Test with nonexistent folder path."""
-        deleted = backup_imap_emails.delete_orphan_local_files(str(tmp_path / "nonexistent"), {"1", "2"})
-        assert deleted == 0
-
-    def test_ignore_non_eml_files(self, tmp_path):
-        """Test that non-.eml files are ignored."""
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-
-        (inbox_path / "1_Email.eml").write_bytes(b"content")
-        (inbox_path / "notes.txt").write_bytes(b"some notes")
-        (inbox_path / "2_data.json").write_bytes(b"{}")
-
-        # Server only has UID 1
-        server_uids = {"1"}
-
-        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
-
-        # Only .eml files should be considered
-        assert deleted == 0
-        assert (inbox_path / "notes.txt").exists()
-        assert (inbox_path / "2_data.json").exists()
-
-
-class TestDestDeleteBackupArgument:
-    """Tests for --dest-delete argument in backup script."""
-
-    def test_dest_delete_default_false(self):
-        """Test that --dest-delete defaults to False."""
-        import argparse
-
-        parser = argparse.ArgumentParser()
-        parser.add_argument("folder", nargs="?")
-        parser.add_argument("--dest-delete", action="store_true", default=False)
-        args = parser.parse_args([])
-        assert args.dest_delete is False
-
-    def test_dest_delete_when_set(self):
-        """Test that --dest-delete can be set to True."""
-        import argparse
-
-        parser = argparse.ArgumentParser()
-        parser.add_argument("folder", nargs="?")
-        parser.add_argument("--dest-delete", action="store_true", default=False)
-        args = parser.parse_args(["--dest-delete"])
-        assert args.dest_delete is True
-
-
-class TestDestDeleteBackupEnvVar:
-    """End-to-end tests for DEST_DELETE env var wiring in backup script."""
-
-    def test_dest_delete_enabled_via_env_var_deletes_orphans(self, single_mock_server, monkeypatch, tmp_path):
-        """If DEST_DELETE=true and server folder is empty, local orphans are deleted."""
-        src_data = {"INBOX": []}
-        _, port = single_mock_server(src_data)
-
-        backup_root = tmp_path / "backup"
-        inbox_path = backup_root / "INBOX"
-        inbox_path.mkdir(parents=True)
-        orphan = inbox_path / "1_Orphan.eml"
-        orphan.write_bytes(b"Subject: Orphan\r\nMessage-ID: <orphan@test>\r\n\r\nBody")
-
-        monkeypatch.setenv("SRC_IMAP_HOST", "localhost")
-        monkeypatch.setenv("SRC_IMAP_USERNAME", "user")
-        monkeypatch.setenv("SRC_IMAP_PASSWORD", "pass")
-        monkeypatch.setenv("BACKUP_LOCAL_PATH", str(backup_root))
-        monkeypatch.setenv("MAX_WORKERS", "1")
-        monkeypatch.setenv("DEST_DELETE", "true")
-
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        backup_imap_emails.main()
-
-        assert not orphan.exists()

test/test_compare_imap_folders.py

@@ -1,425 +0,0 @@
-"""
-Tests for compare_imap_folders.py
-
-Tests cover:
-- Basic folder comparison between accounts
-- Matching counts
-- Mismatched counts
-- Missing folders on destination
-- Empty folder handling
-- Configuration validation
-"""
-
-import imaplib
-import os
-import sys
-from unittest.mock import MagicMock
-
-import pytest
-
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
-
-import compare_imap_folders
-from conftest import make_mock_connection, make_single_mock_connection
-
-
-class TestFolderComparison:
-    """Tests for folder comparison functionality."""
-
-    def test_matching_counts(self, mock_server_factory, monkeypatch, capsys):
-        """Test comparison when source and destination have matching counts."""
-        data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
-            "Sent": [b"Subject: 3\r\n\r\nB"],
-        }
-        _, _, p1, p2 = mock_server_factory(data, data.copy())
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "Sent" in captured.out
-
-    def test_mismatched_counts(self, mock_server_factory, monkeypatch, capsys):
-        """Test comparison when counts differ."""
-        src_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB", b"Subject: 3\r\n\r\nB"],
-        }
-        dest_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB"],
-        }
-        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        # Source has 3, dest has 1
-        assert "3" in captured.out
-        assert "1" in captured.out
-
-    def test_folder_missing_on_destination(self, mock_server_factory, monkeypatch, capsys):
-        """Test when a folder exists on source but not destination."""
-        src_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB"],
-            "Archive": [b"Subject: 2\r\n\r\nB"],
-        }
-        dest_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB"],
-        }
-        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        # Archive should show N/A for destination
-        assert "Archive" in captured.out
-        assert "N/A" in captured.out
-
-
-class TestEmptyFolders:
-    """Tests for empty folder handling."""
-
-    def test_empty_folders(self, mock_server_factory, monkeypatch, capsys):
-        """Test comparison with empty folders."""
-        src_data = {"INBOX": [], "Empty": []}
-        dest_data = {"INBOX": [], "Empty": []}
-
-        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "0" in captured.out
-
-
-class TestGetEmailCount:
-    """Tests for get_email_count function."""
-
-    def test_successful_count(self, single_mock_server):
-        """Test successful email count."""
-        data = {"INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"]}
-        _, port = single_mock_server(data)
-
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        result = compare_imap_folders.get_email_count(conn, "INBOX")
-        assert result == 2
-
-        conn.logout()
-
-    def test_nonexistent_folder(self, single_mock_server):
-        """Test count for non-existent folder returns None."""
-        data = {"INBOX": []}
-        _, port = single_mock_server(data)
-
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        result = compare_imap_folders.get_email_count(conn, "NonExistent")
-        assert result is None
-
-        conn.logout()
-
-
-class TestCompareFoldersErrorHandling:
-    """Tests for error handling in compare_imap_folders.py"""
-
-    def test_connection_failure(self, monkeypatch):
-        """Test graceful exit when connection fails."""
-        mock_get = MagicMock(return_value=None)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-
-        # Test source fail
-        env = {
-            "SRC_IMAP_HOST": "h",
-            "SRC_IMAP_USERNAME": "u",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "h",
-            "DEST_IMAP_USERNAME": "u",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-
-        compare_imap_folders.main()
-        # Should call once and exit
-        assert mock_get.call_count == 1
-
-    def test_dest_connection_failure(self, monkeypatch):
-        """Test graceful exit when destination connection fails."""
-        # Source OK, Dest None
-        mock_src = MagicMock()
-        mock_src.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"'])
-
-        def side_effect(h, u, p):
-            if u == "src_u":
-                return mock_src
-            return None
-
-        monkeypatch.setattr("imap_common.get_imap_connection", side_effect)
-
-        env = {
-            "SRC_IMAP_HOST": "h",
-            "SRC_IMAP_USERNAME": "src_u",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "h",
-            "DEST_IMAP_USERNAME": "dest_u",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-
-        compare_imap_folders.main()
-
-    def test_list_failure(self, monkeypatch, capsys):
-        """Test handling of LIST command failure."""
-        mock_src = MagicMock()
-        mock_src.list.return_value = ("NO", [])
-        mock_dest = MagicMock()
-
-        monkeypatch.setattr(
-            "imap_common.get_imap_connection", lambda h, u, p, oauth2_token=None: mock_src if u == "s" else mock_dest
-        )
-
-        env = {
-            "SRC_IMAP_HOST": "h",
-            "SRC_IMAP_USERNAME": "s",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "h",
-            "DEST_IMAP_USERNAME": "d",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        assert "Failed to list source folders" in captured.out
-
-    def test_select_failure(self, single_mock_server):
-        """Test get_email_count handles select failure."""
-        data = {"INBOX": []}
-        _, port = single_mock_server(data)
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        # Mocking select failure on a real connection object is hard without
-        # using a pure mock. So let's use a Mock object instead of real conn.
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("NO", [b"Error"])
-
-        result = compare_imap_folders.get_email_count(mock_conn, "INBOX")
-        assert result is None
-
-    def test_search_failure(self, single_mock_server):
-        """Test get_email_count handles search failure."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"Selected"])
-        mock_conn.search.return_value = ("NO", [b"Error"])
-
-        result = compare_imap_folders.get_email_count(mock_conn, "INBOX")
-        assert result is None
-
-    def test_imap_exception_in_count(self):
-        """Test exception handling in get_email_count."""
-        mock_conn = MagicMock()
-        mock_conn.select.side_effect = imaplib.IMAP4.error("Crash")
-
-        result = compare_imap_folders.get_email_count(mock_conn, "INBOX")
-        assert result is None
-
-
-class TestConfigValidation:
-    """Tests for configuration validation."""
-
-    def test_missing_source_credentials(self, monkeypatch, capsys):
-        """Test that missing source credentials cause exit."""
-        env = {
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-
-        with pytest.raises(SystemExit) as exc_info:
-            compare_imap_folders.main()
-
-        assert exc_info.value.code == 2
-
-    def test_missing_dest_credentials(self, monkeypatch, capsys):
-        """Test that missing destination credentials cause exit."""
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src",
-            "SRC_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-
-        with pytest.raises(SystemExit) as exc_info:
-            compare_imap_folders.main()
-
-        assert exc_info.value.code == 2
-
-
-class TestTotals:
-    """Tests for total calculation."""
-
-    def test_total_calculation(self, mock_server_factory, monkeypatch, capsys):
-        """Test that totals are calculated correctly."""
-        src_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
-            "Sent": [b"Subject: 3\r\n\r\nB", b"Subject: 4\r\n\r\nB", b"Subject: 5\r\n\r\nB"],
-        }
-        dest_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB"],
-            "Sent": [b"Subject: 3\r\n\r\nB"],
-        }
-        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["compare_imap_folders.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        # Total source: 5, Total dest: 2, Diff: 3
-        assert "TOTAL" in captured.out
-        assert "5" in captured.out
-        assert "2" in captured.out
-
-
-class TestLocalFolderComparison:
-    """Tests for comparing IMAP folders to local backup folders."""
-
-    def test_local_source_to_imap_dest(self, single_mock_server, monkeypatch, tmp_path, capsys):
-        """Local source folder list drives the comparison; destination is IMAP."""
-        # Local source: INBOX has 2 emails, Archive has 1
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
-        (inbox_path / "2_b.eml").write_bytes(b"Subject: B\r\n\r\nBody")
-
-        archive_path = tmp_path / "Archive"
-        archive_path.mkdir()
-        (archive_path / "1_c.eml").write_bytes(b"Subject: C\r\n\r\nBody")
-
-        # Destination IMAP: INBOX has 1, Archive is missing
-        dest_data = {"INBOX": [b"Subject: 1\r\n\r\nB"]}
-        _, port = single_mock_server(dest_data)
-
-        env = {
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["compare_imap_folders.py", "--src-path", str(tmp_path)],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "Archive" in captured.out
-        # Destination should show N/A for missing Archive
-        assert "N/A" in captured.out
-
-    def test_imap_source_to_local_dest(self, single_mock_server, monkeypatch, tmp_path, capsys):
-        """IMAP source folder list drives the comparison; destination is local."""
-        # Source IMAP: INBOX has 2, Sent has 1
-        src_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
-            "Sent": [b"Subject: 3\r\n\r\nB"],
-        }
-        _, port = single_mock_server(src_data)
-
-        # Local dest: INBOX has 1, Sent missing
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["compare_imap_folders.py", "--dest-path", str(tmp_path)],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        compare_imap_folders.main()
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "Sent" in captured.out
-        assert "N/A" in captured.out

test/test_count_imap_emails.py

@@ -1,256 +0,0 @@
-"""
-Tests for count_imap_emails.py
-
-Tests cover:
-- Basic email counting
-- Multiple folder counting
-- Empty folder handling
-- Error handling
-- Configuration validation
-"""
-
-import imaplib
-import os
-import sys
-from unittest.mock import MagicMock
-
-import pytest
-
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
-
-import count_imap_emails
-from conftest import make_single_mock_connection
-
-
-class TestEmailCounting:
-    """Tests for email counting functionality."""
-
-    def test_count_single_folder(self, single_mock_server, monkeypatch, capsys):
-        """Test counting emails in a single folder."""
-        src_data = {
-            "INBOX": [
-                b"Subject: Email 1\r\n\r\nBody",
-                b"Subject: Email 2\r\n\r\nBody",
-                b"Subject: Email 3\r\n\r\nBody",
-            ]
-        }
-        _, port = single_mock_server(src_data)
-
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        count_imap_emails.count_emails("localhost", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "3" in captured.out
-
-    def test_count_multiple_folders(self, single_mock_server, monkeypatch, capsys):
-        """Test counting emails across multiple folders."""
-        src_data = {
-            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
-            "Sent": [b"Subject: 3\r\n\r\nB"],
-            "Archive": [b"Subject: 4\r\n\r\nB", b"Subject: 5\r\n\r\nB", b"Subject: 6\r\n\r\nB"],
-        }
-        _, port = single_mock_server(src_data)
-
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        count_imap_emails.count_emails("localhost", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "Sent" in captured.out
-        assert "Archive" in captured.out
-        # Total should be 6
-        assert "6" in captured.out
-
-    def test_empty_folder(self, single_mock_server, monkeypatch, capsys):
-        """Test counting in empty folders."""
-        src_data = {"INBOX": [], "Empty": []}
-        _, port = single_mock_server(src_data)
-
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        count_imap_emails.count_emails("localhost", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "0" in captured.out
-
-
-class TestLocalEmailCounting:
-    """Tests for counting emails from a local backup folder."""
-
-    def test_count_local_folders(self, tmp_path, capsys):
-        inbox_path = tmp_path / "INBOX"
-        inbox_path.mkdir()
-        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
-        (inbox_path / "2_b.eml").write_bytes(b"Subject: B\r\n\r\nBody")
-
-        gmail_all_mail = tmp_path / "[Gmail]" / "All Mail"
-        gmail_all_mail.mkdir(parents=True)
-        (gmail_all_mail / "1_c.eml").write_bytes(b"Subject: C\r\n\r\nBody")
-
-        count_imap_emails.count_local_emails(str(tmp_path))
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-        assert "[Gmail]/All Mail" in captured.out
-        assert "TOTAL" in captured.out
-        assert "3" in captured.out
-
-
-class TestEmailCountingErrors:
-    """Tests for error handling in email counting."""
-
-    def test_connection_failure(self, monkeypatch):
-        """Test graceful exit when connection fails."""
-        mock_get = MagicMock(return_value=None)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-
-        # Should return silently without raising
-        count_imap_emails.count_emails("host", "user", "pass")
-        mock_get.assert_called_once()
-
-    def test_list_command_failure(self, monkeypatch):
-        """Test handling of LIST command failure."""
-        mock_mail = MagicMock()
-        mock_mail.list.return_value = ("NO", [])
-        mock_get = MagicMock(return_value=mock_mail)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-
-        count_imap_emails.count_emails("host", "user", "pass")
-        mock_mail.list.assert_called_once()
-        # Should exit early, so select should not be called
-        mock_mail.select.assert_not_called()
-
-    def test_select_command_failure(self, monkeypatch):
-        """Test handling of SELECT command failure for a folder."""
-        mock_mail = MagicMock()
-        mock_mail.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"'])
-        # Fail selection
-        mock_mail.select.return_value = ("NO", [b"Select failed"])
-
-        mock_get = MagicMock(return_value=mock_mail)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-        monkeypatch.setattr("imap_common.normalize_folder_name", lambda x: "INBOX")
-
-        count_imap_emails.count_emails("host", "user", "pass")
-
-        mock_mail.select.assert_called_once()
-        # Should skip search for this folder
-        mock_mail.search.assert_not_called()
-
-    def test_search_command_failure(self, monkeypatch, capsys):
-        """Test handling of SEARCH command failure."""
-        mock_mail = MagicMock()
-        mock_mail.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"'])
-        mock_mail.select.return_value = ("OK", [b"Selected"])
-        # Fail search
-        mock_mail.search.return_value = ("NO", [b"Search failed"])
-
-        mock_get = MagicMock(return_value=mock_mail)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-        monkeypatch.setattr("imap_common.normalize_folder_name", lambda x: "INBOX")
-
-        count_imap_emails.count_emails("host", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "Error" in captured.out
-
-    def test_imap_exception_during_list(self, monkeypatch, capsys):
-        """Test handling of IMAP4 exception during list command."""
-        mock_mail = MagicMock()
-        mock_mail.list.side_effect = imaplib.IMAP4.error("Crash listing")
-        mock_get = MagicMock(return_value=mock_mail)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-
-        count_imap_emails.count_emails("host", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "Failed to list mailboxes" in captured.out
-
-    def test_imap_exception_during_select(self, monkeypatch, capsys):
-        """Test handling of IMAP4 exception during folder selection."""
-        mock_mail = MagicMock()
-        mock_mail.list.return_value = ("OK", [rb'(\HasNoChildren) "/" "INBOX"'])
-        mock_mail.select.side_effect = imaplib.IMAP4.error("Crash selecting")
-
-        mock_get = MagicMock(return_value=mock_mail)
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-        monkeypatch.setattr("imap_common.normalize_folder_name", lambda x: "INBOX")
-
-        count_imap_emails.count_emails("host", "user", "pass")
-
-        captured = capsys.readouterr()
-        # Should print Error for that folder
-        assert "Error" in captured.out
-
-    def test_generic_exception(self, monkeypatch, capsys):
-        """Test handling of generic connection/runtime exceptions."""
-        mock_get = MagicMock(side_effect=Exception("Generic Crash"))
-        monkeypatch.setattr("imap_common.get_imap_connection", mock_get)
-
-        count_imap_emails.count_emails("host", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "An error occurred: Generic Crash" in captured.out
-
-
-class TestMainFunction:
-    """Tests for main function and CLI."""
-
-    def test_main_with_env_vars(self, single_mock_server, monkeypatch, capsys):
-        """Test main function with environment variables."""
-        src_data = {"INBOX": [b"Subject: Test\r\n\r\nBody"]}
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "IMAP_HOST": "localhost",
-            "IMAP_USERNAME": "user",
-            "IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["count_imap_emails.py"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        # Import and run __main__ block logic
-        count_imap_emails.count_emails("localhost", "user", "pass")
-
-        captured = capsys.readouterr()
-        assert "INBOX" in captured.out
-
-    def test_missing_credentials(self, monkeypatch, capsys):
-        """Test that missing auth is rejected by argparse (neither password nor OAuth2 client-id)."""
-        monkeypatch.setattr(os, "environ", {})
-        monkeypatch.setattr(sys, "argv", ["count_imap_emails.py", "--host", "localhost", "--user", "user"])
-
-        with pytest.raises(SystemExit) as exc_info:
-            count_imap_emails.main()
-
-        assert exc_info.value.code == 2
-
-
-class TestSrcImapFallback:
-    """Tests for SRC_IMAP_* environment variable fallback."""
-
-    def test_src_imap_vars_fallback(self, single_mock_server, monkeypatch, capsys):
-        """Test that SRC_IMAP_* vars work as fallback."""
-        src_data = {"INBOX": [b"Subject: Test\r\n\r\nBody"]}
-        _, port = single_mock_server(src_data)
-
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "user",
-            "SRC_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        # The fallback logic: IMAP_* or SRC_IMAP_*
-        default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
-        default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
-        default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
-
-        assert default_host == "localhost"
-        assert default_user == "user"
-        assert default_pass == "pass"

test/test_imap_backup.py

@@ -0,0 +1,565 @@
+"""
+Tests for backup_imap_emails.py
+
+Tests cover:
+- Basic email backup to local files
+- Incremental backup (skip existing)
+- Multiple folder backup
+- Filename sanitization
+- Configuration validation
+"""
+
+import imaplib
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_backup as backup_imap_emails
+from conftest import temp_argv, temp_env
+from providers import provider_gmail
+from utils import imap_common
+
+
+def _mock_imap_env(port):
+    return {
+        "SRC_IMAP_HOST": f"imap://localhost:{port}",
+        "SRC_IMAP_USERNAME": "user",
+        "SRC_IMAP_PASSWORD": "pass",
+        "MAX_WORKERS": "1",
+        "BATCH_SIZE": "1",
+    }
+
+
+class TestBackupBasic:
+    """Tests for basic backup functionality."""
+
+    def test_single_email_backup(self, single_mock_server, tmp_path):
+        """Test backing up a single email to local directory."""
+        src_data = {"INBOX": [b"Subject: Test Email\r\nMessage-ID: <1@test>\r\n\r\nBody content"]}
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path)]):
+            backup_imap_emails.main()
+
+        # Check that backup was created
+        inbox_path = tmp_path / "INBOX"
+        assert inbox_path.exists()
+
+        eml_files = list(inbox_path.glob("*.eml"))
+        assert len(eml_files) == 1
+
+        # Check content
+        content = eml_files[0].read_bytes()
+        assert b"Test Email" in content or b"Body content" in content
+
+    def test_multiple_emails_backup(self, single_mock_server, tmp_path):
+        """Test backing up multiple emails."""
+        src_data = {
+            "INBOX": [
+                b"Subject: Email 1\r\nMessage-ID: <1@test>\r\n\r\nBody 1",
+                b"Subject: Email 2\r\nMessage-ID: <2@test>\r\n\r\nBody 2",
+                b"Subject: Email 3\r\nMessage-ID: <3@test>\r\n\r\nBody 3",
+            ]
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path)]):
+            backup_imap_emails.main()
+
+        inbox_path = tmp_path / "INBOX"
+        eml_files = list(inbox_path.glob("*.eml"))
+        assert len(eml_files) == 3
+
+
+class TestIncrementalBackup:
+    """Tests for incremental backup functionality."""
+
+    def test_skip_existing_emails(self, single_mock_server, tmp_path):
+        """Test that existing emails are skipped during incremental backup."""
+        src_data = {
+            "INBOX": [
+                b"Subject: Email 1\r\nMessage-ID: <1@test>\r\n\r\nBody 1",
+                b"Subject: Email 2\r\nMessage-ID: <2@test>\r\n\r\nBody 2",
+            ]
+        }
+        _, port = single_mock_server(src_data)
+
+        # Create pre-existing file
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        existing_file = inbox_path / "1_Email_1.eml"
+        existing_file.write_bytes(b"existing content")
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path)]):
+            backup_imap_emails.main()
+
+        # Should still have original content (not overwritten)
+        assert existing_file.read_bytes() == b"existing content"
+
+        # But second email should be backed up
+        eml_files = list(inbox_path.glob("*.eml"))
+        assert len(eml_files) == 2
+
+
+class TestMultipleFolderBackup:
+    """Tests for backing up multiple folders."""
+
+    def test_backup_all_folders(self, single_mock_server, tmp_path):
+        """Test backing up emails from multiple folders."""
+        src_data = {
+            "INBOX": [b"Subject: Inbox\r\nMessage-ID: <1@test>\r\n\r\nC"],
+            "Sent": [b"Subject: Sent\r\nMessage-ID: <2@test>\r\n\r\nC"],
+            "Archive": [b"Subject: Archive\r\nMessage-ID: <3@test>\r\n\r\nC"],
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path)]):
+            backup_imap_emails.main()
+
+        assert (tmp_path / "INBOX").exists()
+        assert (tmp_path / "Sent").exists()
+        assert (tmp_path / "Archive").exists()
+
+    def test_backup_single_folder(self, single_mock_server, tmp_path):
+        """Test backing up a specific folder only."""
+        src_data = {
+            "INBOX": [b"Subject: Inbox\r\nMessage-ID: <1@test>\r\n\r\nC"],
+            "Sent": [b"Subject: Sent\r\nMessage-ID: <2@test>\r\n\r\nC"],
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path), "INBOX"]):
+            backup_imap_emails.main()
+
+        assert (tmp_path / "INBOX").exists()
+        # Sent should NOT be backed up
+        assert not (tmp_path / "Sent").exists()
+
+
+class TestEmptyFolderHandling:
+    """Tests for empty folder handling."""
+
+    def test_empty_folder(self, single_mock_server, tmp_path):
+        """Test handling of empty folders."""
+        src_data = {"INBOX": [], "Empty": []}
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path)]):
+            backup_imap_emails.main()
+
+
+class TestConfigValidation:
+    """Tests for configuration validation."""
+
+    def test_missing_credentials(self, capsys, tmp_path):
+        """Test that missing credentials cause exit."""
+        with temp_env({}), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path)]):
+            with pytest.raises(SystemExit) as exc_info:
+                backup_imap_emails.main()
+
+        assert exc_info.value.code == 2
+
+    def test_missing_dest_path(self, capsys):
+        """Test that missing destination path causes exit."""
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "user",
+            "SRC_IMAP_PASSWORD": "pass",
+        }
+        with temp_env(env), temp_argv(["backup_imap_emails.py"]):
+            with pytest.raises(SystemExit) as exc_info:
+                backup_imap_emails.main()
+
+        assert exc_info.value.code == 2
+
+
+class TestGetExistingUids:
+    """Tests for get_existing_uids function."""
+
+    def test_empty_directory(self, tmp_path):
+        """Test with empty directory."""
+        result = backup_imap_emails.get_existing_uids(str(tmp_path))
+        assert result == set()
+
+    def test_nonexistent_directory(self):
+        """Test with non-existent directory."""
+        result = backup_imap_emails.get_existing_uids("/nonexistent/path")
+        assert result == set()
+
+    def test_existing_files(self, tmp_path):
+        """Test extraction of UIDs from existing files."""
+        (tmp_path / "1_Subject.eml").touch()
+        (tmp_path / "2_Another.eml").touch()
+        (tmp_path / "100_Long_Subject.eml").touch()
+
+        result = backup_imap_emails.get_existing_uids(str(tmp_path))
+        assert result == {"1", "2", "100"}
+
+    def test_ignore_non_matching_files(self, tmp_path):
+        """Test that non-matching files are ignored."""
+        (tmp_path / "not_an_email.txt").touch()
+        (tmp_path / "random.eml").touch()  # No underscore
+        (tmp_path / "1_Subject.eml").touch()
+
+        result = backup_imap_emails.get_existing_uids(str(tmp_path))
+        assert result == {"1"}
+
+
+class TestGmailLabelsPreservation:
+    """Tests for Gmail labels manifest functionality."""
+
+    def test_is_label_folder_user_labels(self):
+        """Test that user labels are correctly identified."""
+        assert provider_gmail.is_label_folder("Work") is True
+        assert provider_gmail.is_label_folder("Personal") is True
+        assert provider_gmail.is_label_folder("Projects/2024") is True
+        assert provider_gmail.is_label_folder("INBOX") is True
+
+    def test_is_label_folder_gmail_labels(self):
+        """Test that Gmail system labels are correctly identified."""
+        assert provider_gmail.is_label_folder("[Gmail]/Sent Mail") is True
+        assert provider_gmail.is_label_folder("[Gmail]/Starred") is True
+
+    def test_is_label_folder_system_folders(self):
+        """Test that Gmail system folders are excluded."""
+        assert provider_gmail.is_label_folder("[Gmail]/All Mail") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Spam") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Trash") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Drafts") is False
+        assert provider_gmail.is_label_folder("[Gmail]/Bin") is False
+
+    def test_load_labels_manifest_nonexistent(self, tmp_path):
+        """Test loading manifest when file doesn't exist."""
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
+        assert result == {}
+
+    def test_load_labels_manifest_existing(self, tmp_path):
+        """Test loading existing manifest file."""
+        import json
+
+        manifest_data = {
+            "<msg1@test.com>": ["INBOX", "Work"],
+            "<msg2@test.com>": ["Sent Mail", "Personal"],
+        }
+        manifest_path = tmp_path / "labels_manifest.json"
+        manifest_path.write_text(json.dumps(manifest_data))
+
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
+        assert result == manifest_data
+
+    def test_load_labels_manifest_invalid_json(self, tmp_path):
+        """Test loading invalid JSON manifest file."""
+        manifest_path = tmp_path / "labels_manifest.json"
+        manifest_path.write_text("not valid json {{{")
+
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
+        assert result == {}
+
+    def test_get_message_ids_in_folder(self, single_mock_server):
+        src_data = {
+            "INBOX": [
+                b"Subject: A\r\nMessage-ID: <msg1@test.com>\r\n\r\nBody",
+                b"Subject: B\r\nMessage-ID: <msg2@test.com>\r\n\r\nBody",
+            ]
+        }
+        _server, port = single_mock_server(src_data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        result = backup_imap_emails.get_message_ids_in_folder(conn, "INBOX", None)
+
+        assert "<msg1@test.com>" in result
+        assert "<msg2@test.com>" in result
+        conn.logout()
+
+    def test_get_message_info_in_folder_read_status(self, single_mock_server):
+        src_data = {
+            "INBOX": [
+                {"uid": 1, "flags": {"\\Seen"}, "content": b"Message-ID: <msg1@test.com>\r\n"},
+                {"uid": 2, "flags": set(), "content": b"Message-ID: <msg2@test.com>\r\n"},
+                {"uid": 3, "flags": {"\\Seen", "\\Answered"}, "content": b"Message-ID: <msg3@test.com>\r\n"},
+            ]
+        }
+        _server, port = single_mock_server(src_data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        result = backup_imap_emails.get_message_info_in_folder(conn, "INBOX", None)
+
+        assert "<msg1@test.com>" in result
+        assert "\\Seen" in result["<msg1@test.com>"]["flags"]
+        assert "<msg2@test.com>" in result
+        assert result["<msg2@test.com>"]["flags"] == []
+        assert "<msg3@test.com>" in result
+        assert "\\Seen" in result["<msg3@test.com>"]["flags"]
+        assert "\\Answered" in result["<msg3@test.com>"]["flags"]
+        conn.logout()
+
+    def test_build_labels_manifest(self, single_mock_server, tmp_path):
+        src_data = {
+            "INBOX": [b"Message-ID: <msg1@test.com>\r\n"],
+            "Work": [b"Message-ID: <msg1@test.com>\r\n"],
+            "[Gmail]/Sent Mail": [b"Message-ID: <msg2@test.com>\r\n"],
+            "[Gmail]/All Mail": [
+                {"uid": 1, "flags": {"\\Seen", "\\Flagged"}, "content": b"Message-ID: <msg1@test.com>\r\n"},
+                {"uid": 2, "flags": set(), "content": b"Message-ID: <msg2@test.com>\r\n"},
+            ],
+        }
+        _server, port = single_mock_server(src_data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        result = backup_imap_emails.build_labels_manifest(conn, str(tmp_path))
+
+        assert "<msg1@test.com>" in result
+        assert "<msg2@test.com>" in result
+        assert "INBOX" in result["<msg1@test.com>"]["labels"]
+        assert "Work" in result["<msg1@test.com>"]["labels"]
+        assert "\\Seen" in result["<msg1@test.com>"]["flags"]
+        assert "\\Flagged" in result["<msg1@test.com>"]["flags"]
+        assert "Sent Mail" in result["<msg2@test.com>"]["labels"]
+        assert result["<msg2@test.com>"]["flags"] == []
+
+        manifest_path = tmp_path / "labels_manifest.json"
+        assert manifest_path.exists()
+        conn.logout()
+
+    def test_preserve_labels_flag_integration(self, single_mock_server, tmp_path):
+        """Test --preserve-labels flag creates manifest."""
+        src_data = {
+            "INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+            "Work": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+            "[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with (
+            temp_env(env),
+            temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path), "--preserve-labels", "[Gmail]/All Mail"]),
+        ):
+            backup_imap_emails.main()
+
+        # Check manifest was created
+        manifest_path = tmp_path / "labels_manifest.json"
+        assert manifest_path.exists()
+
+    def test_manifest_only_flag(self, single_mock_server, tmp_path):
+        """Test --manifest-only flag creates manifest without downloading emails."""
+        src_data = {
+            "INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+            "Work": [b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody"],
+            "[Gmail]/All Mail": [
+                b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody",
+                b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody",
+            ],
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path), "--manifest-only"]):
+            # Should exit with code 0 after creating manifest
+            with pytest.raises(SystemExit) as exc_info:
+                backup_imap_emails.main()
+
+        assert exc_info.value.code == 0
+
+        # Check manifest was created
+        manifest_path = tmp_path / "labels_manifest.json"
+        assert manifest_path.exists()
+
+        # Check NO email folders were created (no download happened)
+        # Only the manifest file should exist
+        items = list(tmp_path.iterdir())
+        assert len(items) == 1
+        assert items[0].name == "labels_manifest.json"
+
+    def test_gmail_system_folders_constant(self):
+        """Test that GMAIL_SYSTEM_FOLDERS contains expected entries."""
+        expected = {
+            "[Gmail]/All Mail",
+            "[Gmail]/Spam",
+            "[Gmail]/Trash",
+            "[Gmail]/Drafts",
+            "[Gmail]/Bin",
+            "[Gmail]/Important",
+        }
+        assert provider_gmail.GMAIL_SYSTEM_FOLDERS == expected
+
+    def test_gmail_mode_flag(self, single_mock_server, tmp_path):
+        """Test --gmail-mode flag backs up All Mail and creates manifest."""
+        src_data = {
+            "INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+            "Work": [b"Subject: Work Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+            "[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["backup_imap_emails.py", "--dest-path", str(tmp_path), "--gmail-mode"]):
+            backup_imap_emails.main()
+
+        # Check manifest was created
+        manifest_path = tmp_path / "labels_manifest.json"
+        assert manifest_path.exists()
+
+        # Check only [Gmail]/All Mail folder was backed up (not INBOX or Work)
+        all_mail_folder = tmp_path / "[Gmail]" / "All Mail"
+        assert all_mail_folder.exists()
+
+        # Should NOT have created separate INBOX or Work folders
+        inbox_folder = tmp_path / "INBOX"
+        work_folder = tmp_path / "Work"
+        assert not inbox_folder.exists()
+        assert not work_folder.exists()
+
+
+class TestDeleteOrphanLocalFiles:
+    """Tests for delete_orphan_local_files function."""
+
+    def test_delete_orphan_files(self, tmp_path):
+        """Test that local files not on server are deleted."""
+        # Create a local folder with some .eml files
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+
+        # Create files with UIDs 1, 2, 3
+        (inbox_path / "1_Test_Email.eml").write_bytes(b"content1")
+        (inbox_path / "2_Another_Email.eml").write_bytes(b"content2")
+        (inbox_path / "3_Third_Email.eml").write_bytes(b"content3")
+
+        # Server only has UIDs 1 and 3
+        server_uids = {"1", "3"}
+
+        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
+
+        # UID 2 should be deleted
+        assert deleted == 1
+        assert (inbox_path / "1_Test_Email.eml").exists()
+        assert not (inbox_path / "2_Another_Email.eml").exists()
+        assert (inbox_path / "3_Third_Email.eml").exists()
+
+    def test_delete_multiple_orphans(self, tmp_path):
+        """Test deleting multiple orphan files."""
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+
+        # Create files with UIDs 1-5
+        for i in range(1, 6):
+            (inbox_path / f"{i}_Email_{i}.eml").write_bytes(f"content{i}".encode())
+
+        # Server only has UID 3
+        server_uids = {"3"}
+
+        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
+
+        # 4 files should be deleted (UIDs 1, 2, 4, 5)
+        assert deleted == 4
+        assert not (inbox_path / "1_Email_1.eml").exists()
+        assert not (inbox_path / "2_Email_2.eml").exists()
+        assert (inbox_path / "3_Email_3.eml").exists()
+        assert not (inbox_path / "4_Email_4.eml").exists()
+        assert not (inbox_path / "5_Email_5.eml").exists()
+
+    def test_no_orphans(self, tmp_path):
+        """Test when all local files exist on server."""
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+
+        (inbox_path / "1_Email.eml").write_bytes(b"content")
+        (inbox_path / "2_Email.eml").write_bytes(b"content")
+
+        server_uids = {"1", "2"}
+
+        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
+
+        assert deleted == 0
+        assert (inbox_path / "1_Email.eml").exists()
+        assert (inbox_path / "2_Email.eml").exists()
+
+    def test_nonexistent_folder(self, tmp_path):
+        """Test with nonexistent folder path."""
+        deleted = backup_imap_emails.delete_orphan_local_files(str(tmp_path / "nonexistent"), {"1", "2"})
+        assert deleted == 0
+
+    def test_ignore_non_eml_files(self, tmp_path):
+        """Test that non-.eml files are ignored."""
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+
+        (inbox_path / "1_Email.eml").write_bytes(b"content")
+        (inbox_path / "notes.txt").write_bytes(b"some notes")
+        (inbox_path / "2_data.json").write_bytes(b"{}")
+
+        # Server only has UID 1
+        server_uids = {"1"}
+
+        deleted = backup_imap_emails.delete_orphan_local_files(str(inbox_path), server_uids)
+
+        # Only .eml files should be considered
+        assert deleted == 0
+        assert (inbox_path / "notes.txt").exists()
+        assert (inbox_path / "2_data.json").exists()
+
+
+class TestDestDeleteBackupArgument:
+    """Tests for --dest-delete argument in backup script."""
+
+    def test_dest_delete_default_false(self):
+        """Test that --dest-delete defaults to False."""
+        import argparse
+
+        parser = argparse.ArgumentParser()
+        parser.add_argument("folder", nargs="?")
+        parser.add_argument("--dest-delete", action="store_true", default=False)
+        args = parser.parse_args([])
+        assert args.dest_delete is False
+
+    def test_dest_delete_when_set(self):
+        """Test that --dest-delete can be set to True."""
+        import argparse
+
+        parser = argparse.ArgumentParser()
+        parser.add_argument("folder", nargs="?")
+        parser.add_argument("--dest-delete", action="store_true", default=False)
+        args = parser.parse_args(["--dest-delete"])
+        assert args.dest_delete is True
+
+
+class TestDestDeleteBackupEnvVar:
+    """End-to-end tests for DEST_DELETE env var wiring in backup script."""
+
+    def test_dest_delete_enabled_via_env_var_deletes_orphans(self, single_mock_server, tmp_path):
+        """If DEST_DELETE=true and server folder is empty, local orphans are deleted."""
+        src_data = {"INBOX": []}
+        _, port = single_mock_server(src_data)
+
+        backup_root = tmp_path / "backup"
+        inbox_path = backup_root / "INBOX"
+        inbox_path.mkdir(parents=True)
+        orphan = inbox_path / "1_Orphan.eml"
+        orphan.write_bytes(b"Subject: Orphan\r\nMessage-ID: <orphan@test>\r\n\r\nBody")
+
+        env = _mock_imap_env(port)
+        env.update(
+            {
+                "BACKUP_LOCAL_PATH": str(backup_root),
+                "DEST_DELETE": "true",
+            }
+        )
+        with temp_env(env), temp_argv(["backup_imap_emails.py"]):
+            backup_imap_emails.main()
+
+        assert not orphan.exists()

test/test_imap_compare.py

@@ -0,0 +1,257 @@
+"""
+Tests for compare_imap_folders.py
+
+Tests cover:
+- Basic folder comparison between accounts
+- Matching counts
+- Mismatched counts
+- Missing folders on destination
+- Empty folder handling
+- Configuration validation
+"""
+
+import imaplib
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_compare as compare_imap_folders
+from conftest import temp_argv, temp_env
+
+
+def _mock_compare_env(src_port, dest_port):
+    return {
+        "SRC_IMAP_HOST": f"imap://localhost:{src_port}",
+        "SRC_IMAP_USERNAME": "src_user",
+        "SRC_IMAP_PASSWORD": "p",
+        "DEST_IMAP_HOST": f"imap://localhost:{dest_port}",
+        "DEST_IMAP_USERNAME": "dest_user",
+        "DEST_IMAP_PASSWORD": "p",
+    }
+
+
+class TestFolderComparison:
+    """Tests for folder comparison functionality."""
+
+    def test_matching_counts(self, mock_server_factory, capsys):
+        """Test comparison when source and destination have matching counts."""
+        data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
+            "Sent": [b"Subject: 3\r\n\r\nB"],
+        }
+        _, _, p1, p2 = mock_server_factory(data, data.copy())
+
+        env = _mock_compare_env(p1, p2)
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Sent" in captured.out
+
+    def test_mismatched_counts(self, mock_server_factory, capsys):
+        """Test comparison when counts differ."""
+        src_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB", b"Subject: 3\r\n\r\nB"],
+        }
+        dest_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB"],
+        }
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        env = _mock_compare_env(p1, p2)
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        # Source has 3, dest has 1
+        assert "3" in captured.out
+        assert "1" in captured.out
+
+    def test_folder_missing_on_destination(self, mock_server_factory, capsys):
+        """Test when a folder exists on source but not destination."""
+        src_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB"],
+            "Archive": [b"Subject: 2\r\n\r\nB"],
+        }
+        dest_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB"],
+        }
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        env = _mock_compare_env(p1, p2)
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        # Archive should show N/A for destination
+        assert "Archive" in captured.out
+        assert "N/A" in captured.out
+
+
+class TestEmptyFolders:
+    """Tests for empty folder handling."""
+
+    def test_empty_folders(self, mock_server_factory, capsys):
+        """Test comparison with empty folders."""
+        src_data = {"INBOX": [], "Empty": []}
+        dest_data = {"INBOX": [], "Empty": []}
+
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        env = _mock_compare_env(p1, p2)
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "0" in captured.out
+
+
+class TestGetEmailCount:
+    """Tests for get_email_count function."""
+
+    def test_successful_count(self, single_mock_server):
+        """Test successful email count."""
+        data = {"INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"]}
+        _, port = single_mock_server(data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        result = compare_imap_folders.get_email_count(conn, "INBOX")
+        assert result == 2
+
+        conn.logout()
+
+    def test_nonexistent_folder(self, single_mock_server):
+        """Test count for non-existent folder returns None."""
+        data = {"INBOX": []}
+        _, port = single_mock_server(data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        result = compare_imap_folders.get_email_count(conn, "NonExistent")
+        assert result is None
+
+        conn.logout()
+
+
+class TestConfigValidation:
+    """Tests for configuration validation."""
+
+    def test_missing_source_credentials(self, capsys):
+        """Test that missing source credentials cause exit."""
+        env = {
+            "DEST_IMAP_HOST": "localhost",
+            "DEST_IMAP_USERNAME": "dest",
+            "DEST_IMAP_PASSWORD": "p",
+        }
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            with pytest.raises(SystemExit) as exc_info:
+                compare_imap_folders.main()
+
+        assert exc_info.value.code == 2
+
+    def test_missing_dest_credentials(self, capsys):
+        """Test that missing destination credentials cause exit."""
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "src",
+            "SRC_IMAP_PASSWORD": "p",
+        }
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            with pytest.raises(SystemExit) as exc_info:
+                compare_imap_folders.main()
+
+        assert exc_info.value.code == 2
+
+
+class TestTotals:
+    """Tests for total calculation."""
+
+    def test_total_calculation(self, mock_server_factory, capsys):
+        """Test that totals are calculated correctly."""
+        src_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
+            "Sent": [b"Subject: 3\r\n\r\nB", b"Subject: 4\r\n\r\nB", b"Subject: 5\r\n\r\nB"],
+        }
+        dest_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB"],
+            "Sent": [b"Subject: 3\r\n\r\nB"],
+        }
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        env = _mock_compare_env(p1, p2)
+        with temp_env(env), temp_argv(["compare_imap_folders.py"]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        # Total source: 5, Total dest: 2, Diff: 3
+        assert "TOTAL" in captured.out
+        assert "5" in captured.out
+        assert "2" in captured.out
+
+
+class TestLocalFolderComparison:
+    """Tests for comparing IMAP folders to local backup folders."""
+
+    def test_local_source_to_imap_dest(self, single_mock_server, tmp_path, capsys):
+        """Local source folder list drives the comparison; destination is IMAP."""
+        # Local source: INBOX has 2 emails, Archive has 1
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+        (inbox_path / "2_b.eml").write_bytes(b"Subject: B\r\n\r\nBody")
+
+        archive_path = tmp_path / "Archive"
+        archive_path.mkdir()
+        (archive_path / "1_c.eml").write_bytes(b"Subject: C\r\n\r\nBody")
+
+        # Destination IMAP: INBOX has 1, Archive is missing
+        dest_data = {"INBOX": [b"Subject: 1\r\n\r\nB"]}
+        _, port = single_mock_server(dest_data)
+
+        env = {
+            "DEST_IMAP_HOST": f"imap://localhost:{port}",
+            "DEST_IMAP_USERNAME": "dest_user",
+            "DEST_IMAP_PASSWORD": "p",
+        }
+        with temp_env(env), temp_argv(["compare_imap_folders.py", "--src-path", str(tmp_path)]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Archive" in captured.out
+        # Destination should show N/A for missing Archive
+        assert "N/A" in captured.out
+
+    def test_imap_source_to_local_dest(self, single_mock_server, tmp_path, capsys):
+        """IMAP source folder list drives the comparison; destination is local."""
+        # Source IMAP: INBOX has 2, Sent has 1
+        src_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
+            "Sent": [b"Subject: 3\r\n\r\nB"],
+        }
+        _, port = single_mock_server(src_data)
+
+        # Local dest: INBOX has 1, Sent missing
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+        env = {
+            "SRC_IMAP_HOST": f"imap://localhost:{port}",
+            "SRC_IMAP_USERNAME": "src_user",
+            "SRC_IMAP_PASSWORD": "p",
+        }
+        with temp_env(env), temp_argv(["compare_imap_folders.py", "--dest-path", str(tmp_path)]):
+            compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Sent" in captured.out
+        assert "N/A" in captured.out

test/test_imap_count.py

@@ -0,0 +1,302 @@
+"""
+Tests for count_imap_emails.py
+
+Tests cover:
+- Basic email counting
+- Multiple folder counting
+- Empty folder handling
+- Error handling
+- Configuration validation
+"""
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_count as count_imap_emails
+from conftest import temp_argv, temp_env
+from utils import imap_common
+
+
+def _mock_imap_env(port):
+    return {
+        "IMAP_HOST": f"imap://localhost:{port}",
+        "IMAP_USERNAME": "user",
+        "IMAP_PASSWORD": "pass",
+    }
+
+
+class TestEmailCounting:
+    """Tests for email counting functionality."""
+
+    def test_count_single_folder(self, single_mock_server, capsys):
+        """Test counting emails in a single folder."""
+        src_data = {
+            "INBOX": [
+                b"Subject: Email 1\r\n\r\nBody",
+                b"Subject: Email 2\r\n\r\nBody",
+                b"Subject: Email 3\r\n\r\nBody",
+            ]
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env):
+            count_imap_emails.count_emails(env["IMAP_HOST"], env["IMAP_USERNAME"], env["IMAP_PASSWORD"])
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "3" in captured.out
+
+    def test_count_multiple_folders(self, single_mock_server, capsys):
+        """Test counting emails across multiple folders."""
+        src_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
+            "Sent": [b"Subject: 3\r\n\r\nB"],
+            "Archive": [b"Subject: 4\r\n\r\nB", b"Subject: 5\r\n\r\nB", b"Subject: 6\r\n\r\nB"],
+        }
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env):
+            count_imap_emails.count_emails(env["IMAP_HOST"], env["IMAP_USERNAME"], env["IMAP_PASSWORD"])
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Sent" in captured.out
+        assert "Archive" in captured.out
+        # Total should be 6
+        assert "6" in captured.out
+
+    def test_empty_folder(self, single_mock_server, capsys):
+        """Test counting in empty folders."""
+        src_data = {"INBOX": [], "Empty": []}
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env):
+            count_imap_emails.count_emails(env["IMAP_HOST"], env["IMAP_USERNAME"], env["IMAP_PASSWORD"])
+
+        captured = capsys.readouterr()
+        assert "0" in captured.out
+
+
+class TestLocalEmailCounting:
+    """Tests for counting emails from a local backup folder."""
+
+    def test_count_local_folders(self, tmp_path, capsys):
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+        (inbox_path / "2_b.eml").write_bytes(b"Subject: B\r\n\r\nBody")
+
+        gmail_all_mail = tmp_path / "[Gmail]" / "All Mail"
+        gmail_all_mail.mkdir(parents=True)
+        (gmail_all_mail / "1_c.eml").write_bytes(b"Subject: C\r\n\r\nBody")
+
+        count_imap_emails.count_local_emails(str(tmp_path))
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "[Gmail]/All Mail" in captured.out
+        assert "TOTAL" in captured.out
+        assert "3" in captured.out
+
+    def test_count_local_ignores_hidden_dirs(self, tmp_path, capsys):
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+        (inbox_path / "note.txt").write_text("ignore")
+
+        hidden_path = tmp_path / ".hidden"
+        hidden_path.mkdir()
+        (hidden_path / "1_hidden.eml").write_bytes(b"Subject: Hidden\r\n\r\nBody")
+
+        cache_path = tmp_path / "__pycache__"
+        cache_path.mkdir()
+        (cache_path / "1_cache.eml").write_bytes(b"Subject: Cache\r\n\r\nBody")
+
+        nested_path = tmp_path / "Projects" / "Sub"
+        nested_path.mkdir(parents=True)
+        (nested_path / "1_sub.eml").write_bytes(b"Subject: Sub\r\n\r\nBody")
+
+        count_imap_emails.count_local_emails(str(tmp_path))
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Projects/Sub" in captured.out
+        assert ".hidden" not in captured.out
+        assert "__pycache__" not in captured.out
+
+    def test_get_local_email_count_unreadable_folder(self, tmp_path):
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+
+        os.chmod(inbox_path, 0)
+        try:
+            result = imap_common.get_local_email_count(str(tmp_path), "INBOX")
+            assert result is None
+        finally:
+            os.chmod(inbox_path, 0o700)
+
+
+class TestImapCommonHelpers:
+    """Tests for imap_common helpers via script tests."""
+
+    def test_list_selectable_folders_filters_noselect(self):
+        class FakeConn:
+            def list(self):
+                return (
+                    "OK",
+                    [
+                        b'(\\Noselect) "/" "Archive"',
+                        b'(\\HasNoChildren) "/" "INBOX"',
+                        '(\\HasNoChildren) "/" "Sent"',
+                    ],
+                )
+
+        result = imap_common.list_selectable_folders(FakeConn())
+        assert result == ["INBOX", "Sent"]
+
+    def test_list_selectable_folders_list_error(self):
+        class FakeConn:
+            def list(self):
+                return ("NO", [])
+
+        result = imap_common.list_selectable_folders(FakeConn())
+        assert result == []
+
+    def test_list_selectable_folders_exception(self):
+        class FakeConn:
+            def list(self):
+                raise Exception("list failed")
+
+        result = imap_common.list_selectable_folders(FakeConn())
+        assert result == []
+
+    def test_get_imap_connection_oauth2_uses_authenticate(self):
+        from unittest.mock import patch
+
+        class FakeIMAP:
+            def __init__(self, _host):
+                self.auth_called = False
+                self.login_called = False
+
+            def authenticate(self, _mechanism, auth_cb):
+                self.auth_called = True
+                auth_cb(None)
+
+            def login(self, _user, _password):
+                self.login_called = True
+
+        with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP):
+            conn = imap_common.get_imap_connection("host", "user", oauth2_token="token")
+
+        assert conn.auth_called is True
+        assert conn.login_called is False
+
+    def test_get_imap_connection_basic_login(self):
+        from unittest.mock import patch
+
+        class FakeIMAP:
+            def __init__(self, _host):
+                self.auth_called = False
+                self.login_called = False
+
+            def authenticate(self, _mechanism, _auth_cb):
+                self.auth_called = True
+
+            def login(self, _user, _password):
+                self.login_called = True
+
+        with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP):
+            conn = imap_common.get_imap_connection("host", "user", password="pass")
+
+        assert conn.login_called is True
+        assert conn.auth_called is False
+
+    def test_ensure_connection_returns_same_conn_when_healthy(self):
+        class GoodConn:
+            def __init__(self):
+                self.noop_calls = 0
+
+            def noop(self):
+                self.noop_calls += 1
+
+        conn = GoodConn()
+        result = imap_common.ensure_connection(conn, "host", "user", "pass")
+        assert result is conn
+        assert conn.noop_calls == 1
+
+    def test_ensure_connection_reconnects_on_noop_error(self):
+        from unittest.mock import patch
+
+        class BadConn:
+            def noop(self):
+                raise Exception("fail")
+
+        new_conn = object()
+        with patch.object(imap_common, "get_imap_connection", return_value=new_conn):
+            result = imap_common.ensure_connection(BadConn(), "host", "user", "pass")
+        assert result is new_conn
+
+    def test_ensure_connection_from_conf_reconnects_on_noop_error(self):
+        from unittest.mock import patch
+
+        class BadConn:
+            def noop(self):
+                raise Exception("fail")
+
+        new_conn = object()
+        with patch.object(imap_common, "get_imap_connection_from_conf", return_value=new_conn):
+            result = imap_common.ensure_connection_from_conf(BadConn(), {"host": "h", "user": "u"})
+        assert result is new_conn
+
+
+class TestMainFunction:
+    """Tests for main function and CLI."""
+
+    def test_main_with_env_vars(self, single_mock_server, capsys):
+        """Test main function with environment variables."""
+        src_data = {"INBOX": [b"Subject: Test\r\n\r\nBody"]}
+        _, port = single_mock_server(src_data)
+
+        env = _mock_imap_env(port)
+        with temp_env(env), temp_argv(["count_imap_emails.py"]):
+            count_imap_emails.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+
+    def test_missing_credentials(self, capsys):
+        """Test that missing auth is rejected by argparse (neither password nor OAuth2 client-id)."""
+        with temp_env({}), temp_argv(["count_imap_emails.py", "--host", "localhost", "--user", "user"]):
+            with pytest.raises(SystemExit) as exc_info:
+                count_imap_emails.main()
+
+        assert exc_info.value.code == 2
+
+
+class TestSrcImapFallback:
+    """Tests for SRC_IMAP_* environment variable fallback."""
+
+    def test_src_imap_vars_fallback(self, capsys):
+        """Test that SRC_IMAP_* vars work as fallback."""
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "user",
+            "SRC_IMAP_PASSWORD": "pass",
+        }
+        with temp_env(env):
+            # The fallback logic: IMAP_* or SRC_IMAP_*
+            default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
+            default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
+            default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
+
+            assert default_host == "localhost"
+            assert default_user == "user"
+            assert default_pass == "pass"

test/test_imap_migrate.py

@@ -19,38 +19,42 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import migrate_imap_emails
-from conftest import make_mock_connection
+import imap_migrate as migrate_imap_emails
+from conftest import temp_argv, temp_env
+from utils import imap_common
+
+
+def _mock_migrate_env(src_port, dest_port):
+    return {
+        "SRC_IMAP_HOST": f"imap://localhost:{src_port}",
+        "SRC_IMAP_USERNAME": "src_user",
+        "SRC_IMAP_PASSWORD": "p",
+        "DEST_IMAP_HOST": f"imap://localhost:{dest_port}",
+        "DEST_IMAP_USERNAME": "dest_user",
+        "DEST_IMAP_PASSWORD": "p",
+        "MAX_WORKERS": "1",
+        "BATCH_SIZE": "1",
+    }
 
 
 class TestBasicMigration:
     """Tests for basic migration functionality."""
 
-    def test_single_email_migration(self, mock_server_factory, monkeypatch):
+    def test_single_email_migration(self, mock_server_factory):
         """Test migrating a single email from source to destination."""
         src_data = {"INBOX": [b"Subject: Hello\r\nMessage-ID: <1@test>\r\n\r\nBody"]}
         dest_data = {"INBOX": []}
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert b"Subject: Hello" in dest_server.folders["INBOX"][0]["content"]
 
-    def test_multiple_emails_migration(self, mock_server_factory, monkeypatch):
+    def test_multiple_emails_migration(self, mock_server_factory):
         """Test migrating multiple emails."""
         src_data = {
             "INBOX": [
@@ -63,19 +67,9 @@ class TestBasicMigration:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 3
 
@@ -83,7 +77,7 @@ class TestBasicMigration:
 class TestDuplicateHandling:
     """Tests for duplicate detection and skipping."""
 
-    def test_skip_duplicate_by_message_id(self, mock_server_factory, monkeypatch):
+    def test_skip_duplicate_by_message_id(self, mock_server_factory):
         """Test that emails with existing Message-ID are skipped."""
         msg = b"Subject: Dup\r\nMessage-ID: <dup-id>\r\n\r\nContent"
         src_data = {"INBOX": [msg]}
@@ -91,24 +85,14 @@ class TestDuplicateHandling:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         # Should still be 1 message (duplicate skipped)
         assert len(dest_server.folders["INBOX"]) == 1
 
-    def test_migrate_non_duplicate(self, mock_server_factory, monkeypatch):
+    def test_migrate_non_duplicate(self, mock_server_factory):
         """Test that non-duplicate emails are migrated even when others exist."""
         existing = b"Subject: Existing\r\nMessage-ID: <existing>\r\n\r\nOld"
         new_msg = b"Subject: New\r\nMessage-ID: <new>\r\n\r\nNew content"
@@ -118,19 +102,9 @@ class TestDuplicateHandling:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         # Should now have 2 messages
         assert len(dest_server.folders["INBOX"]) == 2
@@ -139,52 +113,31 @@ class TestDuplicateHandling:
 class TestDeleteFromSource:
     """Tests for delete-after-migration functionality."""
 
-    def test_delete_after_migration(self, mock_server_factory, monkeypatch):
+    def test_delete_after_migration(self, mock_server_factory):
         """Test that emails are deleted from source after successful migration."""
         src_data = {"INBOX": [b"Subject: Del\r\nMessage-ID: <del>\r\n\r\nC"]}
         dest_data = {"INBOX": []}
 
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-            "DELETE_FROM_SOURCE": "true",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["DELETE_FROM_SOURCE"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert len(src_server.folders["INBOX"]) == 0
 
-    def test_no_delete_when_disabled(self, mock_server_factory, monkeypatch):
+    def test_no_delete_when_disabled(self, mock_server_factory):
         """Test that emails remain in source when delete is not enabled."""
         src_data = {"INBOX": [b"Subject: Keep\r\nMessage-ID: <keep>\r\n\r\nC"]}
         dest_data = {"INBOX": []}
 
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-            # DELETE_FROM_SOURCE not set
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert len(src_server.folders["INBOX"]) == 1
@@ -193,31 +146,21 @@ class TestDeleteFromSource:
 class TestFolderHandling:
     """Tests for folder creation and multi-folder migration."""
 
-    def test_folder_creation(self, mock_server_factory, monkeypatch):
+    def test_folder_creation(self, mock_server_factory):
         """Test that folders are created on destination."""
         src_data = {"INBOX": [], "Archive": [b"Subject: A\r\nMessage-ID: <a>\r\n\r\nC"]}
         dest_data = {"INBOX": []}
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
+            migrate_imap_emails.main()
 
         assert "Archive" in dest_server.folders
         assert len(dest_server.folders["Archive"]) == 1
 
-    def test_multiple_folders_migration(self, mock_server_factory, monkeypatch):
+    def test_multiple_folders_migration(self, mock_server_factory):
         """Test migrating emails from multiple folders."""
         src_data = {
             "INBOX": [b"Subject: Inbox\r\nMessage-ID: <inbox>\r\n\r\nC"],
@@ -228,49 +171,28 @@ class TestFolderHandling:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert "Sent" in dest_server.folders
         assert "Drafts" in dest_server.folders
 
-    def test_empty_folder_handling(self, mock_server_factory, monkeypatch):
+    def test_empty_folder_handling(self, mock_server_factory):
         """Test that empty folders are handled gracefully."""
         src_data = {"INBOX": [], "Empty": []}
         dest_data = {"INBOX": []}
 
         _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        # Should complete without error
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
+            migrate_imap_emails.main()
 
 
 class TestPreserveFlags:
-    def test_preserve_flags_applied_on_copy(self, mock_server_factory, monkeypatch):
+    def test_preserve_flags_applied_on_copy(self, mock_server_factory):
         msg = b"Subject: Flagged\r\nMessage-ID: <f1@test>\r\n\r\nBody"
         src_data = {"INBOX": [msg]}
         dest_data = {"INBOX": []}
@@ -278,25 +200,15 @@ class TestPreserveFlags:
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
         src_server.folders["INBOX"][0]["flags"].add("\\Seen")
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-            "PRESERVE_FLAGS": "true",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["PRESERVE_FLAGS"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert "\\Seen" in dest_server.folders["INBOX"][0]["flags"]
 
-    def test_preserve_flags_syncs_on_duplicate(self, mock_server_factory, monkeypatch):
+    def test_preserve_flags_syncs_on_duplicate(self, mock_server_factory):
         msg = b"Subject: DupFlags\r\nMessage-ID: <df1@test>\r\n\r\nBody"
         src_data = {"INBOX": [msg]}
         dest_data = {"INBOX": [msg]}
@@ -304,27 +216,17 @@ class TestPreserveFlags:
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
         src_server.folders["INBOX"][0]["flags"].add("\\Seen")
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-            "PRESERVE_FLAGS": "true",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["PRESERVE_FLAGS"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert "\\Seen" in dest_server.folders["INBOX"][0]["flags"]
 
 
 class TestGmailModeLabels:
-    def test_gmail_mode_migrates_all_mail_and_applies_labels(self, mock_server_factory, monkeypatch):
+    def test_gmail_mode_migrates_all_mail_and_applies_labels(self, mock_server_factory):
         msg1 = b"Subject: One\r\nMessage-ID: <gm1@test>\r\n\r\nBody 1"
         msg2 = b"Subject: Two\r\nMessage-ID: <gm2@test>\r\n\r\nBody 2"
 
@@ -337,54 +239,108 @@ class TestGmailModeLabels:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-            "GMAIL_MODE": "true",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["GMAIL_MODE"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
+            migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert "Work" in dest_server.folders
         assert len(dest_server.folders["Work"]) == 2
 
+    def test_gmail_mode_fallback_folder_for_unlabeled(self, mock_server_factory):
+        msg = b"Subject: Unlabeled\r\nMessage-ID: <gm-unlabeled@test>\r\n\r\nBody"
+
+        src_data = {
+            "[Gmail]/All Mail": [msg],
+        }
+        dest_data = {"INBOX": []}
+
+        _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        env = _mock_migrate_env(p1, p2)
+        env["GMAIL_MODE"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
+            migrate_imap_emails.main()
+
+        assert imap_common.FOLDER_RESTORED_UNLABELED in dest_server.folders
+        assert len(dest_server.folders[imap_common.FOLDER_RESTORED_UNLABELED]) == 1
+
+
+class TestCacheHitWithoutLock:
+    """Covers cached skip when no lock is provided."""
+
+    def test_cached_skip_without_lock(self, mock_server_factory):
+        msg_id = "<cache-no-lock@test>"
+        msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": []}
+
+        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        src = imaplib.IMAP4("localhost", p1)
+        dest = imaplib.IMAP4("localhost", p2)
+        src.login("src_user", "p")
+        dest.login("dest_user", "p")
+
+        src.select('"INBOX"', readonly=False)
+
+        existing_dest_msg_ids_by_folder = {"INBOX": {msg_id}}
+        success, _src, _dest, deleted = migrate_imap_emails.process_single_uid(
+            src,
+            dest,
+            b"1",
+            "INBOX",
+            False,
+            None,
+            False,
+            False,
+            None,
+            True,
+            False,
+            existing_dest_msg_ids_by_folder=existing_dest_msg_ids_by_folder,
+            existing_dest_msg_ids_lock=None,
+            progress_cache_path=None,
+            progress_cache_data=None,
+            progress_cache_lock=None,
+            dest_host="localhost",
+            dest_user="dest_user",
+        )
+
+        assert success is True
+        assert deleted == 0
+        assert len(dest_server.folders["INBOX"]) == 0
+
+        src.logout()
+        dest.logout()
+
 
 class TestConfigValidation:
     """Tests for configuration validation."""
 
-    def test_missing_source_credentials(self, monkeypatch, capsys):
+    def test_missing_source_credentials(self, capsys):
         """Test that missing source credentials cause exit."""
         env = {
             "DEST_IMAP_HOST": "localhost",
             "DEST_IMAP_USERNAME": "dest",
             "DEST_IMAP_PASSWORD": "p",
         }
-        monkeypatch.setattr(os, "environ", env)
-
-        with pytest.raises(SystemExit) as exc_info:
-            migrate_imap_emails.main()
+        with temp_env(env):
+            with pytest.raises(SystemExit) as exc_info:
+                migrate_imap_emails.main()
 
         assert exc_info.value.code == 2
 
-    def test_missing_dest_credentials(self, monkeypatch, capsys):
+    def test_missing_dest_credentials(self, capsys):
         """Test that missing destination credentials cause exit."""
         env = {
             "SRC_IMAP_HOST": "localhost",
             "SRC_IMAP_USERNAME": "src",
             "SRC_IMAP_PASSWORD": "p",
         }
-        monkeypatch.setattr(os, "environ", env)
-
-        with pytest.raises(SystemExit) as exc_info:
-            migrate_imap_emails.main()
+        with temp_env(env):
+            with pytest.raises(SystemExit) as exc_info:
+                migrate_imap_emails.main()
 
         assert exc_info.value.code == 2
 
@@ -392,38 +348,32 @@ class TestConfigValidation:
 class TestMigrateErrorHandling:
     """Tests for error handling during migration."""
 
-    def test_connection_error_in_worker(self, mock_server_factory, monkeypatch):
+    def test_connection_error_in_worker(self, mock_server_factory):
         """Test error handling when worker fails to connect."""
         src_data = {"INBOX": [b"Subject: Test\r\nMessage-ID: <1>\r\n\r\nBody"]}
         dest_data = {"INBOX": []}
 
         _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        # Proper setup for main thread connection
-        real_mock_conn = make_mock_connection(p1, p2)
+        from unittest.mock import patch
 
-        def side_effect(*args, **kwargs):
+        original_get_imap_connection = imap_common.get_imap_connection
+
+        def side_effect(host, user, pwd, oauth2_token=None):
             if threading.current_thread() is threading.main_thread():
-                return real_mock_conn(*args, **kwargs)
+                return original_get_imap_connection(host, user, pwd, oauth2_token)
             return None  # Simulate connection failure in worker
 
-        monkeypatch.setattr("imap_common.get_imap_connection", side_effect)
+        env = _mock_migrate_env(p1, p2)
+        with (
+            temp_env(env),
+            temp_argv(["migrate_imap_emails.py", "INBOX"]),
+            patch("utils.imap_common.get_imap_connection", side_effect=side_effect),
+        ):
+            # Should not crash, but log error
+            migrate_imap_emails.main()
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-
-        # Should not crash, but log error
-        migrate_imap_emails.main()
-
-    def test_select_error_in_worker(self, mock_server_factory, monkeypatch):
+    def test_select_error_in_worker(self, mock_server_factory):
         """Test error handling when folder selection fails in worker."""
         src_data = {"INBOX": [b"Subject: Test\r\nMessage-ID: <1>\r\n\r\nBody"]}
         dest_data = {"INBOX": []}
@@ -438,23 +388,44 @@ class TestMigrateErrorHandling:
                 raise RuntimeError("Select failed")
             return original_select(self, mailbox, readonly)
 
-        monkeypatch.setattr(imaplib.IMAP4, "select", side_effect_select)
+        from unittest.mock import patch
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+        env = _mock_migrate_env(p1, p2)
+        with (
+            patch.object(imaplib.IMAP4, "select", side_effect=side_effect_select),
+            temp_env(env),
+            temp_argv(["migrate_imap_emails.py", "INBOX"]),
+        ):
+            migrate_imap_emails.main()
 
-        migrate_imap_emails.main()
+    def test_select_error_in_process_batch(self, mock_server_factory):
+        """Cover process_batch select exception handling with real server data."""
+        src_data = {"INBOX": [b"Subject: Test\r\nMessage-ID: <1>\r\n\r\nBody"]}
+        dest_data = {"INBOX": []}
 
-    def test_fetch_error_in_worker(self, mock_server_factory, monkeypatch):
+        _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        def raise_select(_self, _mailbox, readonly=False):
+            raise RuntimeError("Select failed")
+
+        from unittest.mock import patch
+
+        env = _mock_migrate_env(p1, p2)
+        src_conf = {"host": env["SRC_IMAP_HOST"], "user": "src_user", "password": "p"}
+        dest_conf = {"host": env["DEST_IMAP_HOST"], "user": "dest_user", "password": "p"}
+
+        with patch.object(imaplib.IMAP4, "select", raise_select), temp_env(env):
+            migrate_imap_emails.process_batch(
+                [b"1"],
+                "INBOX",
+                src_conf,
+                dest_conf,
+                delete_from_source=False,
+                preserve_flags=False,
+                gmail_mode=False,
+            )
+
+    def test_fetch_error_in_worker(self, mock_server_factory):
         """Test error handling when fetching message details fails."""
         src_data = {"INBOX": [b"Subject: Test\r\nMessage-ID: <1>\r\n\r\nBody"]}
         dest_data = {"INBOX": []}
@@ -469,28 +440,22 @@ class TestMigrateErrorHandling:
                 raise RuntimeError("Fetch failed")
             return original_uid(self, command, *args)
 
-        monkeypatch.setattr(imaplib.IMAP4, "uid", side_effect_uid)
+        from unittest.mock import patch
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with (
+            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid, autospec=True),
+            temp_env(env),
+            temp_argv(["migrate_imap_emails.py", "INBOX"]),
+        ):
+            migrate_imap_emails.main()
 
         # Dest should be empty as fetch failed
         assert len(dest_server.folders["INBOX"]) == 0
 
-    def test_main_connection_failure(self, monkeypatch):
+    def test_main_connection_failure(self):
         """Test that main exits if initial connection fails."""
-        monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: None)
+        from unittest.mock import patch
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -500,45 +465,60 @@ class TestMigrateErrorHandling:
             "DEST_IMAP_USERNAME": "dest_user",
             "DEST_IMAP_PASSWORD": "p",
         }
-        monkeypatch.setattr(os, "environ", env)
-
-        with pytest.raises(SystemExit) as exc:
-            migrate_imap_emails.main()
+        with patch("utils.imap_common.get_imap_connection", return_value=None), temp_env(env):
+            with pytest.raises(SystemExit) as exc:
+                migrate_imap_emails.main()
         assert exc.value.code == 1
 
+    def test_main_logs_progress_cache_load_failure(self, mock_server_factory, capsys):
+        """Cover progress cache load exception in main."""
+        src_data = {"INBOX": [b"Subject: X\r\nMessage-ID: <x@test>\r\n\r\nBody"]}
+        dest_data = {"INBOX": []}
+
+        _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        from unittest.mock import patch
+
+        env = _mock_migrate_env(p1, p2)
+        with (
+            patch(
+                "utils.imap_common.load_progress_cache",
+                lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")),
+            ),
+            temp_env(env),
+            temp_argv(["migrate_imap_emails.py", "--migrate-cache", "./cache"]),
+        ):
+            migrate_imap_emails.main()
+
+        captured = capsys.readouterr()
+        assert "Warning: Failed to load progress cache" in captured.out
+
 
 class TestTrashHandling:
     """Tests for trash folder related logic."""
 
-    def test_circular_trash_migration_prevention(self, mock_server_factory, monkeypatch):
+    def test_circular_trash_migration_prevention(self, mock_server_factory):
         """Test that the trash folder itself is not migrated when delete is on."""
         src_data = {"INBOX": [], "Trash": [b"Subject: Garbage\r\nMessage-ID: <g>\r\n\r\nC"]}
         dest_data = {"INBOX": []}
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        # Mock trash detection
-        monkeypatch.setattr("imap_common.detect_trash_folder", lambda conn: "Trash")
+        from unittest.mock import patch
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "DELETE_FROM_SOURCE": "true",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["DELETE_FROM_SOURCE"] = "true"
+        with (
+            patch("utils.imap_common.detect_trash_folder", return_value="Trash"),
+            temp_env(env),
+            temp_argv(["migrate_imap_emails.py"]),
+        ):
+            migrate_imap_emails.main()
 
         # Trash folder should NOT be created in dest (skipped)
         assert "Trash" not in dest_server.folders
 
-    def test_deleted_moved_to_trash(self, mock_server_factory, monkeypatch):
+    def test_deleted_moved_to_trash(self, mock_server_factory):
         """Test that migrated emails are moved to trash on source."""
         src_data = {"INBOX": [b"Subject: T\r\nMessage-ID: <1>\r\n\r\nBody"]}
         dest_data = {"INBOX": []}
@@ -546,22 +526,16 @@ class TestTrashHandling:
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
         src_server.folders["Trash"] = []  # Create trash on source
 
-        monkeypatch.setattr("imap_common.detect_trash_folder", lambda conn: "Trash")
+        from unittest.mock import patch
 
-        env = {
-            "SRC_IMAP_HOST": "localhost",
-            "SRC_IMAP_USERNAME": "src_user",
-            "SRC_IMAP_PASSWORD": "p",
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "dest_user",
-            "DEST_IMAP_PASSWORD": "p",
-            "DELETE_FROM_SOURCE": "true",
-            "MAX_WORKERS": "1",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["DELETE_FROM_SOURCE"] = "true"
+        with (
+            patch("utils.imap_common.detect_trash_folder", return_value="Trash"),
+            temp_env(env),
+            temp_argv(["migrate_imap_emails.py", "INBOX"]),
+        ):
+            migrate_imap_emails.main()
 
         # Dest should have it
         assert len(dest_server.folders["INBOX"]) == 1
@@ -571,6 +545,58 @@ class TestTrashHandling:
         assert len(src_server.folders["Trash"]) == 1
 
 
+class TestCommonMessageParsing:
+    """Covers imap_common message parsing helpers used by migrate."""
+
+    def test_parse_message_id_from_empty_bytes(self):
+        assert imap_common.parse_message_id_from_bytes(b"") is None
+
+    def test_parse_message_id_and_subject_from_empty_bytes(self):
+        msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(b"")
+        assert msg_id is None
+        assert subject == "(No Subject)"
+
+    def test_get_uid_to_message_id_map_empty(self):
+        result = imap_common.get_uid_to_message_id_map(object(), [])
+        assert result == {}
+
+    def test_extract_message_id_invalid_type(self):
+        assert imap_common.extract_message_id(123) is None
+
+    def test_parse_message_id_from_invalid_type(self):
+        assert imap_common.parse_message_id_from_bytes(123) is None
+
+    def test_parse_message_id_from_bytes_success(self):
+        raw_message = b"Subject: X\r\nMessage-ID: <ok@test>\r\n\r\nBody"
+        assert imap_common.parse_message_id_from_bytes(raw_message) == "<ok@test>"
+
+    def test_parse_message_id_and_subject_from_invalid_type(self):
+        msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(123)
+        assert msg_id is None
+        assert subject == "(No Subject)"
+
+    def test_get_uid_to_message_id_map_missing_uid(self):
+        class FakeConn:
+            def uid(self, _cmd, _uids, _opts):
+                return (
+                    "OK",
+                    [
+                        (
+                            b"1 (BODY[HEADER.FIELDS (MESSAGE-ID)] {40}",
+                            b"Message-ID: <x@test>\r\n",
+                        ),
+                        b")",
+                    ],
+                )
+
+        result = imap_common.get_uid_to_message_id_map(FakeConn(), [b"1"])
+        assert result == {}
+
+    def test_decode_mime_header_exception_path(self):
+        result = imap_common.decode_mime_header(["not", "a", "header"])
+        assert result == "['not', 'a', 'header']"
+
+
 class TestFilterPreservableFlags:
     """Tests for filter_preservable_flags function."""
 
@@ -649,7 +675,7 @@ class TestDestDeleteArgument:
 class TestDestDeleteFunctionality:
     """Tests for --dest-delete actual deletion behavior."""
 
-    def test_delete_orphan_emails_removes_extra_dest_emails(self, mock_server_factory, monkeypatch):
+    def test_delete_orphan_emails_removes_extra_dest_emails(self, mock_server_factory):
         """Test that emails in dest but not in source are deleted with --dest-delete."""
         # Source has only 1 email
         src_data = {"INBOX": [b"Subject: Keep Me\r\nMessage-ID: <keep@test>\r\n\r\nBody"]}
@@ -664,24 +690,17 @@ class TestDestDeleteFunctionality:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        monkeypatch.setenv("SRC_IMAP_HOST", "localhost")
-        monkeypatch.setenv("SRC_IMAP_USERNAME", "src_user")
-        monkeypatch.setenv("SRC_IMAP_PASSWORD", "p")
-        monkeypatch.setenv("DEST_IMAP_HOST", "localhost")
-        monkeypatch.setenv("DEST_IMAP_USERNAME", "dest_user")
-        monkeypatch.setenv("DEST_IMAP_PASSWORD", "p")
-        monkeypatch.setenv("MAX_WORKERS", "1")
-        monkeypatch.setenv("DEST_DELETE", "true")
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["DEST_DELETE"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         # Only the email that exists in source should remain
         assert len(dest_server.folders["INBOX"]) == 1
         remaining_content = dest_server.folders["INBOX"][0]["content"]
         assert b"Message-ID: <keep@test>" in remaining_content
 
-    def test_dest_delete_disabled_keeps_extra_emails(self, mock_server_factory, monkeypatch):
+    def test_dest_delete_disabled_keeps_extra_emails(self, mock_server_factory):
         """Test that without --dest-delete, extra dest emails are kept."""
         src_data = {"INBOX": [b"Subject: Source Email\r\nMessage-ID: <src@test>\r\n\r\nBody"]}
         dest_data = {
@@ -692,22 +711,14 @@ class TestDestDeleteFunctionality:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        monkeypatch.delenv("DEST_DELETE", raising=False)
-        monkeypatch.setenv("SRC_IMAP_HOST", "localhost")
-        monkeypatch.setenv("SRC_IMAP_USERNAME", "src_user")
-        monkeypatch.setenv("SRC_IMAP_PASSWORD", "p")
-        monkeypatch.setenv("DEST_IMAP_HOST", "localhost")
-        monkeypatch.setenv("DEST_IMAP_USERNAME", "dest_user")
-        monkeypatch.setenv("DEST_IMAP_PASSWORD", "p")
-        monkeypatch.setenv("MAX_WORKERS", "1")
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         # Both emails should exist (source was copied, dest-only was kept)
         assert len(dest_server.folders["INBOX"]) == 2
 
-    def test_dest_delete_empty_source_deletes_all(self, mock_server_factory, monkeypatch):
+    def test_dest_delete_empty_source_deletes_all(self, mock_server_factory):
         """Test that if source folder is empty, all dest emails are deleted."""
         src_data = {"INBOX": []}
         dest_data = {
@@ -719,17 +730,46 @@ class TestDestDeleteFunctionality:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        monkeypatch.setenv("SRC_IMAP_HOST", "localhost")
-        monkeypatch.setenv("SRC_IMAP_USERNAME", "src_user")
-        monkeypatch.setenv("SRC_IMAP_PASSWORD", "p")
-        monkeypatch.setenv("DEST_IMAP_HOST", "localhost")
-        monkeypatch.setenv("DEST_IMAP_USERNAME", "dest_user")
-        monkeypatch.setenv("DEST_IMAP_PASSWORD", "p")
-        monkeypatch.setenv("MAX_WORKERS", "1")
-        monkeypatch.setenv("DEST_DELETE", "true")
-        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
-
-        migrate_imap_emails.main()
+        env = _mock_migrate_env(p1, p2)
+        env["DEST_DELETE"] = "true"
+        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
+            migrate_imap_emails.main()
 
         # All dest emails should be deleted
         assert len(dest_server.folders["INBOX"]) == 0
+
+    def test_dest_delete_syncs_after_migration(self, mock_server_factory):
+        """End-to-end: delete orphans after a successful migration batch."""
+        src_data = {"INBOX": [b"Subject: Keep\r\nMessage-ID: <keep@test>\r\n\r\nBody"]}
+        dest_data = {
+            "INBOX": [
+                b"Subject: Keep\r\nMessage-ID: <keep@test>\r\n\r\nBody",
+                b"Subject: Orphan\r\nMessage-ID: <orphan@test>\r\n\r\nBody",
+            ]
+        }
+
+        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        src = imaplib.IMAP4("localhost", p1)
+        dest = imaplib.IMAP4("localhost", p2)
+        src.login("src_user", "p")
+        dest.login("dest_user", "p")
+
+        migrate_imap_emails.MAX_WORKERS = 1
+        migrate_imap_emails.BATCH_SIZE = 1
+
+        migrate_imap_emails.migrate_folder(
+            src,
+            dest,
+            "INBOX",
+            False,
+            {"host": "localhost", "user": "src_user", "password": "p"},
+            {"host": "localhost", "user": "dest_user", "password": "p"},
+            dest_delete=True,
+        )
+
+        assert len(dest_server.folders["INBOX"]) == 1
+        assert b"Message-ID: <keep@test>" in dest_server.folders["INBOX"][0]["content"]
+
+        src.logout()
+        dest.logout()

test/test_imap_oauth2.py

@@ -1,489 +0,0 @@
-"""
-Tests for imap_oauth2.py
-
-Tests cover:
-- OAuth2 provider detection
-- Microsoft tenant discovery
-- Microsoft token acquisition and caching
-- Google token acquisition and caching
-- Provider dispatch
-- Token refresh (caching and silent refresh)
-- Thread-safe token refresh
-"""
-
-import json
-import os
-import sys
-from unittest.mock import MagicMock, patch
-
-import pytest
-
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
-
-import imap_oauth2
-
-
-@pytest.fixture(autouse=True)
-def clear_oauth2_caches():
-    """Clear module-level OAuth2 caches between tests."""
-    imap_oauth2._msal_app_cache.clear()
-    imap_oauth2._google_creds_cache.clear()
-    yield
-    imap_oauth2._msal_app_cache.clear()
-    imap_oauth2._google_creds_cache.clear()
-
-
-class TestDetectOauth2Provider:
-    """Tests for detect_oauth2_provider function."""
-
-    def test_microsoft_outlook(self):
-        """Test detects Microsoft from outlook host."""
-        assert imap_oauth2.detect_oauth2_provider("outlook.office365.com") == "microsoft"
-
-    def test_microsoft_office365(self):
-        """Test detects Microsoft from office365 host."""
-        assert imap_oauth2.detect_oauth2_provider("imap.office365.com") == "microsoft"
-
-    def test_microsoft_mixed_case(self):
-        """Test detects Microsoft case-insensitively."""
-        assert imap_oauth2.detect_oauth2_provider("Outlook.Office365.COM") == "microsoft"
-
-    def test_google_gmail(self):
-        """Test detects Google from gmail host."""
-        assert imap_oauth2.detect_oauth2_provider("imap.gmail.com") == "google"
-
-    def test_google_googlemail(self):
-        """Test detects Google from google host."""
-        assert imap_oauth2.detect_oauth2_provider("imap.google.com") == "google"
-
-    def test_unknown_provider(self):
-        """Test returns None for unrecognized host."""
-        assert imap_oauth2.detect_oauth2_provider("imap.example.com") is None
-
-    def test_unknown_yahoo(self):
-        """Test returns None for Yahoo host."""
-        assert imap_oauth2.detect_oauth2_provider("imap.mail.yahoo.com") is None
-
-
-class TestDiscoverMicrosoftTenant:
-    """Tests for discover_microsoft_tenant function."""
-
-    def test_successful_discovery(self):
-        """Test successful tenant ID extraction from OpenID config."""
-        tenant_id = "12345678-abcd-ef01-2345-67890abcdef0"
-        data = {
-            "issuer": f"https://sts.windows.net/{tenant_id}/",
-            "authorization_endpoint": f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/authorize",
-        }
-
-        with patch.object(imap_oauth2, "_fetch_json_https", return_value=data):
-            result = imap_oauth2.discover_microsoft_tenant("user@contoso.com")
-
-        assert result == tenant_id
-
-    def test_domain_extraction(self):
-        """Test that domain is correctly extracted from email."""
-        data = {"issuer": "https://sts.windows.net/abcdef01-2345-6789-abcd-ef0123456789/"}
-
-        with patch.object(imap_oauth2, "_fetch_json_https", return_value=data) as mock_fetch:
-            imap_oauth2.discover_microsoft_tenant("user@example.org")
-            host, path = mock_fetch.call_args[0][0], mock_fetch.call_args[0][1]
-            assert host == "login.microsoftonline.com"
-            assert "example.org" in path
-
-    def test_network_error(self, capsys):
-        """Test returns None on network error."""
-        with patch.object(imap_oauth2, "_fetch_json_https", side_effect=OSError("Connection refused")):
-            result = imap_oauth2.discover_microsoft_tenant("user@invalid.example")
-
-        assert result is None
-        captured = capsys.readouterr()
-        assert "Could not discover" in captured.out
-
-    def test_invalid_json(self, capsys):
-        """Test returns None on invalid JSON response."""
-        with patch.object(
-            imap_oauth2,
-            "_fetch_json_https",
-            side_effect=json.JSONDecodeError("Expecting value", "not json", 0),
-        ):
-            result = imap_oauth2.discover_microsoft_tenant("user@test.com")
-
-        assert result is None
-
-    def test_no_tenant_in_issuer(self, capsys):
-        """Test returns None when issuer has no tenant GUID."""
-        data = {"issuer": "https://sts.windows.net/not-a-guid/"}
-
-        with patch.object(imap_oauth2, "_fetch_json_https", return_value=data):
-            result = imap_oauth2.discover_microsoft_tenant("user@test.com")
-
-        assert result is None
-        captured = capsys.readouterr()
-        assert "Could not extract tenant ID" in captured.out
-
-
-class TestAcquireMicrosoftOauth2Token:
-    """Tests for acquire_microsoft_oauth2_token function."""
-
-    def test_successful_token(self):
-        """Test successful token acquisition with auto-discovery."""
-        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
-            mock_msal = MagicMock()
-            mock_app = MagicMock()
-            mock_app.get_accounts.return_value = []
-            mock_app.initiate_device_flow.return_value = {"user_code": "ABC123", "message": "Go to..."}
-            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "test_token"}
-            mock_msal.PublicClientApplication.return_value = mock_app
-
-            with patch.dict("sys.modules", {"msal": mock_msal}):
-                result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
-
-            assert result == "test_token"
-
-    def test_tenant_discovery_failure(self, capsys):
-        """Test returns None when tenant discovery fails."""
-        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value=None):
-            result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
-
-        assert result is None
-
-    def test_cached_token(self):
-        """Test returns cached token when available."""
-        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
-            mock_msal = MagicMock()
-            mock_app = MagicMock()
-            mock_account = {"username": "user@test.com"}
-            mock_app.get_accounts.return_value = [mock_account]
-            mock_app.acquire_token_silent.return_value = {"access_token": "cached_token"}
-            mock_msal.PublicClientApplication.return_value = mock_app
-
-            with patch.dict("sys.modules", {"msal": mock_msal}):
-                result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
-
-            assert result == "cached_token"
-
-
-class TestAcquireGoogleOauth2Token:
-    """Tests for acquire_google_oauth2_token function."""
-
-    def test_successful_token(self):
-        """Test successful Google token acquisition."""
-        mock_credentials = MagicMock()
-        mock_credentials.token = "google_test_token"
-
-        mock_flow = MagicMock()
-        mock_flow.run_local_server.return_value = mock_credentials
-
-        mock_installed_app_flow = MagicMock()
-        mock_installed_app_flow.from_client_config.return_value = mock_flow
-
-        mock_module = MagicMock()
-        mock_module.InstalledAppFlow = mock_installed_app_flow
-
-        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
-            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
-
-        assert result == "google_test_token"
-
-    def test_missing_library(self):
-        """Test exits when google-auth-oauthlib is not installed."""
-        with patch.dict("sys.modules", {"google_auth_oauthlib": None, "google_auth_oauthlib.flow": None}):
-            with pytest.raises(SystemExit):
-                imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
-
-    def test_no_token_returned(self):
-        """Test returns None when credentials have no token."""
-        mock_credentials = MagicMock()
-        mock_credentials.token = None
-
-        mock_flow = MagicMock()
-        mock_flow.run_local_server.return_value = mock_credentials
-
-        mock_installed_app_flow = MagicMock()
-        mock_installed_app_flow.from_client_config.return_value = mock_flow
-
-        mock_module = MagicMock()
-        mock_module.InstalledAppFlow = mock_installed_app_flow
-
-        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
-            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
-
-        assert result is None
-
-
-class TestAcquireOauth2TokenForProvider:
-    """Tests for acquire_oauth2_token_for_provider dispatch function."""
-
-    def test_dispatch_to_microsoft(self):
-        """Test dispatches to Microsoft when provider is 'microsoft'."""
-        with patch.object(imap_oauth2, "acquire_microsoft_oauth2_token", return_value="ms_token") as mock_ms:
-            result = imap_oauth2.acquire_oauth2_token_for_provider("microsoft", "cid", "user@test.com")
-
-        assert result == "ms_token"
-        mock_ms.assert_called_once_with("cid", "user@test.com")
-
-    def test_dispatch_to_google(self):
-        """Test dispatches to Google when provider is 'google'."""
-        with patch.object(imap_oauth2, "acquire_google_oauth2_token", return_value="g_token") as mock_g:
-            result = imap_oauth2.acquire_oauth2_token_for_provider("google", "cid", "user@gmail.com", "secret")
-
-        assert result == "g_token"
-        mock_g.assert_called_once_with("cid", "secret")
-
-    def test_google_requires_client_secret(self, capsys):
-        """Test returns None when Google is selected without client_secret."""
-        result = imap_oauth2.acquire_oauth2_token_for_provider("google", "cid", "user@gmail.com")
-
-        assert result is None
-        captured = capsys.readouterr()
-        assert "--oauth2-client-secret" in captured.out
-
-    def test_unknown_provider(self, capsys):
-        """Test returns None for unknown provider."""
-        result = imap_oauth2.acquire_oauth2_token_for_provider("yahoo", "cid", "user@yahoo.com")
-
-        assert result is None
-        captured = capsys.readouterr()
-        assert "Unknown OAuth2 provider" in captured.out
-
-
-class TestMicrosoftTokenRefresh:
-    """Tests for Microsoft OAuth2 token caching and refresh."""
-
-    def test_msal_app_cached_on_first_call(self):
-        """Test MSAL app is cached after first call."""
-        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
-            mock_msal = MagicMock()
-            mock_app = MagicMock()
-            mock_app.get_accounts.return_value = []
-            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
-            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token1"}
-            mock_msal.PublicClientApplication.return_value = mock_app
-
-            with patch.dict("sys.modules", {"msal": mock_msal}):
-                imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
-
-            assert ("client-id", "tenant-123") in imap_oauth2._msal_app_cache
-
-    def test_cached_app_reused_on_second_call(self):
-        """Test second call reuses cached MSAL app instead of creating new one."""
-        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
-            mock_msal = MagicMock()
-            mock_app = MagicMock()
-            mock_app.get_accounts.return_value = []
-            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
-            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token1"}
-            mock_msal.PublicClientApplication.return_value = mock_app
-
-            with patch.dict("sys.modules", {"msal": mock_msal}):
-                imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
-
-                # Second call — simulate cached token available (refresh token worked)
-                mock_account = {"username": "user@test.com"}
-                mock_app.get_accounts.return_value = [mock_account]
-                mock_app.acquire_token_silent.return_value = {"access_token": "refreshed_token"}
-
-                result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
-
-            assert result == "refreshed_token"
-            # PublicClientApplication should only have been called once (first call)
-            assert mock_msal.PublicClientApplication.call_count == 1
-
-
-class TestGoogleTokenRefresh:
-    """Tests for Google OAuth2 token caching and refresh."""
-
-    def test_credentials_cached_on_first_call(self):
-        """Test Google credentials are cached after first call."""
-        mock_credentials = MagicMock()
-        mock_credentials.token = "google_token"
-
-        mock_flow = MagicMock()
-        mock_flow.run_local_server.return_value = mock_credentials
-
-        mock_installed_app_flow = MagicMock()
-        mock_installed_app_flow.from_client_config.return_value = mock_flow
-
-        mock_module = MagicMock()
-        mock_module.InstalledAppFlow = mock_installed_app_flow
-
-        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
-            imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
-
-        assert ("client-id", "client-secret") in imap_oauth2._google_creds_cache
-
-    def test_cached_credentials_refreshed_on_second_call(self):
-        """Test second call refreshes cached credentials without opening browser."""
-        # Pre-populate cache with credentials that have a refresh token
-        mock_creds = MagicMock()
-        mock_creds.refresh_token = "refresh_tok"
-        mock_creds.token = "refreshed_google_token"
-        imap_oauth2._google_creds_cache[("client-id", "client-secret")] = mock_creds
-
-        mock_request_module = MagicMock()
-        with patch.dict(
-            "sys.modules",
-            {
-                "google": MagicMock(),
-                "google.auth": MagicMock(),
-                "google.auth.transport": MagicMock(),
-                "google.auth.transport.requests": mock_request_module,
-            },
-        ):
-            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
-
-        assert result == "refreshed_google_token"
-        # Verify refresh was called
-        mock_creds.refresh.assert_called_once()
-
-    def test_falls_back_to_browser_if_refresh_fails(self):
-        """Test falls back to full auth flow if cached token refresh fails."""
-        # Pre-populate cache with credentials whose refresh fails
-        mock_creds = MagicMock()
-        mock_creds.refresh_token = "refresh_tok"
-        mock_creds.refresh.side_effect = Exception("Refresh failed")
-        imap_oauth2._google_creds_cache[("client-id", "client-secret")] = mock_creds
-
-        # Set up the full auth flow
-        mock_credentials = MagicMock()
-        mock_credentials.token = "new_browser_token"
-
-        mock_flow = MagicMock()
-        mock_flow.run_local_server.return_value = mock_credentials
-
-        mock_installed_app_flow = MagicMock()
-        mock_installed_app_flow.from_client_config.return_value = mock_flow
-
-        mock_module = MagicMock()
-        mock_module.InstalledAppFlow = mock_installed_app_flow
-
-        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
-            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
-
-        assert result == "new_browser_token"
-
-
-class TestRefreshOauth2Token:
-    """Tests for thread-safe refresh_oauth2_token function."""
-
-    def test_refreshes_token_and_updates_conf(self):
-        """Test that a new token is acquired and conf["oauth2_token"] is updated."""
-        conf = {
-            "host": "host",
-            "user": "user",
-            "password": "pass",
-            "oauth2_token": "old_token",
-            "oauth2": {
-                "provider": "microsoft",
-                "client_id": "client-id",
-                "email": "user@test.com",
-                "client_secret": None,
-            },
-        }
-
-        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="new_token") as mock_acquire:
-            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
-
-        assert result == "new_token"
-        assert conf["oauth2_token"] == "new_token"
-        mock_acquire.assert_called_once_with("microsoft", "client-id", "user@test.com", None)
-
-    def test_skips_refresh_when_token_already_updated(self):
-        """Test that refresh is skipped if another thread already updated the token."""
-        conf = {
-            "host": "host",
-            "user": "user",
-            "password": "pass",
-            "oauth2_token": "already_refreshed_token",
-            "oauth2": {
-                "provider": "microsoft",
-                "client_id": "client-id",
-                "email": "user@test.com",
-                "client_secret": None,
-            },
-        }
-
-        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider") as mock_acquire:
-            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
-
-        assert result == "already_refreshed_token"
-        assert conf["oauth2_token"] == "already_refreshed_token"
-        mock_acquire.assert_not_called()
-
-    def test_returns_none_on_refresh_failure(self):
-        """Test returns None and leaves conf unchanged when refresh fails."""
-        conf = {
-            "host": "host",
-            "user": "user",
-            "password": "pass",
-            "oauth2_token": "old_token",
-            "oauth2": {
-                "provider": "google",
-                "client_id": "client-id",
-                "email": "user@gmail.com",
-                "client_secret": "secret",
-            },
-        }
-
-        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
-            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
-
-        assert result is None
-        assert conf["oauth2_token"] == "old_token"
-
-    def test_returns_none_when_no_oauth2_context(self):
-        """Test returns None when oauth2 context is not set."""
-        conf = {
-            "host": "host",
-            "user": "user",
-            "password": "pass",
-            "oauth2_token": "old_token",
-            "oauth2": None,
-        }
-
-        result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
-
-        assert result is None
-        assert conf["oauth2_token"] == "old_token"
-
-    def test_concurrent_threads_only_one_refreshes(self):
-        """Test that only one thread performs the refresh when multiple threads compete."""
-        import threading
-        import time
-
-        conf = {
-            "host": "host",
-            "user": "user",
-            "password": "pass",
-            "oauth2_token": "expired_token",
-            "oauth2": {
-                "provider": "microsoft",
-                "client_id": "client-id",
-                "email": "user@test.com",
-                "client_secret": None,
-            },
-        }
-        call_count = {"value": 0}
-        barrier = threading.Barrier(3)  # 3 threads
-
-        def slow_acquire(provider, client_id, email, client_secret):
-            call_count["value"] += 1
-            time.sleep(0.05)  # Simulate network delay
-            return "fresh_token"
-
-        def thread_func():
-            barrier.wait()  # Ensure all threads start at the same time
-            imap_oauth2.refresh_oauth2_token(conf, "expired_token")
-
-        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", side_effect=slow_acquire):
-            threads = [threading.Thread(target=thread_func) for _ in range(3)]
-            for t in threads:
-                t.start()
-            for t in threads:
-                t.join()
-
-        # Only one thread should have called acquire (the first to get the lock).
-        # The other two should see conf["oauth2_token"] changed and skip.
-        assert call_count["value"] == 1
-        assert conf["oauth2_token"] == "fresh_token"

test/test_imap_restore.py

@@ -9,19 +9,28 @@ Tests cover:
 - Configuration validation
 """
 
+import imaplib
 import json
 import os
 import sys
-from unittest.mock import MagicMock
 
 import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_common
-import restore_cache
-import restore_imap_emails
-from conftest import make_single_mock_connection
+import imap_restore as restore_imap_emails
+from conftest import temp_argv, temp_env
+from utils import imap_common, restore_cache
+
+
+def _mock_restore_env(port):
+    return {
+        "DEST_IMAP_HOST": f"imap://localhost:{port}",
+        "DEST_IMAP_USERNAME": "user",
+        "DEST_IMAP_PASSWORD": "pass",
+        "MAX_WORKERS": "1",
+        "BATCH_SIZE": "1",
+    }
 
 
 class TestLoadLabelsManifest:
@@ -36,7 +45,7 @@ class TestLoadLabelsManifest:
         manifest_path = tmp_path / "labels_manifest.json"
         manifest_path.write_text(json.dumps(manifest_data))
 
-        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
         assert result == manifest_data
 
     def test_load_manifest_new_format(self, tmp_path):
@@ -48,14 +57,14 @@ class TestLoadLabelsManifest:
         manifest_path = tmp_path / "labels_manifest.json"
         manifest_path.write_text(json.dumps(manifest_data))
 
-        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
         assert result == manifest_data
         assert "\\Seen" in result["<msg1@test.com>"]["flags"]
         assert result["<msg2@test.com>"]["flags"] == []
 
     def test_load_nonexistent_manifest(self, tmp_path):
         """Test loading when manifest doesn't exist."""
-        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
         assert result == {}
 
     def test_load_invalid_json_manifest(self, tmp_path):
@@ -63,7 +72,7 @@ class TestLoadLabelsManifest:
         manifest_path = tmp_path / "labels_manifest.json"
         manifest_path.write_text("not valid json {{{")
 
-        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
         assert result == {}
 
 
@@ -147,6 +156,15 @@ Body content.
         assert message_id is None
         assert raw_content is None
 
+    def test_parse_eml_file_unknown_charset_subject(self, tmp_path):
+        """Test parsing with a subject that uses an unknown charset."""
+        file_path = tmp_path / "unknown_charset.eml"
+        file_path.write_text("Subject: =?X-UNKNOWN?B?SGVsbG8=?=\r\nMessage-ID: <unknown@test>\r\n\r\nBody")
+
+        message_id, _date_str, _raw_content, subject = restore_imap_emails.parse_eml_file(str(file_path))
+        assert message_id == "<unknown@test>"
+        assert "Hello" in subject
+
 
 class TestGetEmlFiles:
     """Tests for getting .eml files from a folder."""
@@ -175,168 +193,48 @@ class TestGetEmlFiles:
         assert result == []
 
 
-class TestGetBackupFolders:
-    """Tests for scanning backup folder structure."""
-
-    def test_get_backup_folders(self, tmp_path):
-        """Test scanning backup folders."""
-        # Create folder structure
-        inbox = tmp_path / "INBOX"
-        inbox.mkdir()
-        (inbox / "email1.eml").write_text("content")
-
-        sent = tmp_path / "Sent"
-        sent.mkdir()
-        (sent / "email2.eml").write_text("content")
-
-        result = restore_imap_emails.get_backup_folders(str(tmp_path))
-
-        assert len(result) == 2
-        folder_names = [f[0] for f in result]
-        assert "INBOX" in folder_names
-        assert "Sent" in folder_names
-
-    def test_get_backup_folders_nested(self, tmp_path):
-        """Test scanning nested folder structure."""
-        gmail = tmp_path / "[Gmail]"
-        gmail.mkdir()
-        all_mail = gmail / "All Mail"
-        all_mail.mkdir()
-        (all_mail / "email.eml").write_text("content")
-
-        result = restore_imap_emails.get_backup_folders(str(tmp_path))
-
-        assert len(result) == 1
-        assert result[0][0] == "[Gmail]/All Mail"
-
-    def test_get_backup_folders_empty(self, tmp_path):
-        """Test scanning empty backup folder."""
-        result = restore_imap_emails.get_backup_folders(str(tmp_path))
-        assert result == []
-
-
 class TestConfigValidation:
     """Tests for configuration validation."""
 
-    def test_missing_credentials(self, monkeypatch, capsys, tmp_path):
+    def test_missing_credentials(self, capsys, tmp_path):
         """Test that missing credentials cause exit."""
-        env = {}
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "--src-path", str(tmp_path)])
-
-        with pytest.raises(SystemExit) as exc_info:
-            restore_imap_emails.main()
+        with temp_env({}), temp_argv(["restore_imap_emails.py", "--src-path", str(tmp_path)]):
+            with pytest.raises(SystemExit) as exc_info:
+                restore_imap_emails.main()
 
         assert exc_info.value.code == 2
 
-    def test_missing_src_path(self, monkeypatch, capsys):
+    def test_missing_src_path(self, capsys):
         """Test that missing source path causes exit."""
         env = {
             "DEST_IMAP_HOST": "localhost",
             "DEST_IMAP_USERNAME": "user",
             "DEST_IMAP_PASSWORD": "pass",
         }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py"])
-
-        with pytest.raises(SystemExit) as exc_info:
-            restore_imap_emails.main()
+        with temp_env(env), temp_argv(["restore_imap_emails.py"]):
+            with pytest.raises(SystemExit) as exc_info:
+                restore_imap_emails.main()
 
         assert exc_info.value.code == 2
 
-    def test_nonexistent_src_path(self, monkeypatch, capsys):
+    def test_nonexistent_src_path(self, capsys):
         """Test that non-existent source path causes exit."""
         env = {
             "DEST_IMAP_HOST": "localhost",
             "DEST_IMAP_USERNAME": "user",
             "DEST_IMAP_PASSWORD": "pass",
         }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["restore_imap_emails.py", "--src-path", "/nonexistent/path"],
-        )
-
-        with pytest.raises(SystemExit) as exc_info:
-            restore_imap_emails.main()
+        with temp_env(env), temp_argv(["restore_imap_emails.py", "--src-path", "/nonexistent/path"]):
+            with pytest.raises(SystemExit) as exc_info:
+                restore_imap_emails.main()
 
         assert exc_info.value.code == 1
 
 
-class TestUploadEmail:
-    """Tests for email upload functionality."""
-
-    def test_upload_email_success(self, monkeypatch):
-        """Test successful email upload."""
-        mock_conn = MagicMock()
-        mock_conn.create.return_value = ("OK", [])
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.append.return_value = ("OK", [])
-
-        # Mock email_exists_in_folder to return False (not a duplicate)
-        monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False)
-
-        result = restore_imap_emails.upload_email(
-            mock_conn,
-            "INBOX",
-            b"raw email content",
-            '"15-Jan-2024 10:30:00 +0000"',
-            "<test@test.com>",
-        )
-
-        assert result is True
-        mock_conn.append.assert_called_once()
-
-    def test_upload_email_duplicate(self, monkeypatch):
-        """Test upload returns False when message exists."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"1"])
-
-        # 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,
-            "INBOX",
-            b"raw email content",
-            '"15-Jan-2024 10:30:00 +0000"',
-            "<test@test.com>",
-            check_duplicate=True,
-        )
-
-        assert result is False
-        mock_conn.append.assert_not_called()
-
-    def test_upload_email_with_seen_flag(self, monkeypatch):
-        """Test upload with \\Seen flag for read emails."""
-        mock_conn = MagicMock()
-        mock_conn.create.return_value = ("OK", [])
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.append.return_value = ("OK", [])
-
-        # Mock email_exists_in_folder to return False (not a duplicate)
-        monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False)
-
-        result = restore_imap_emails.upload_email(
-            mock_conn,
-            "INBOX",
-            b"raw email content",
-            '"15-Jan-2024 10:30:00 +0000"',
-            "<test@test.com>",
-            flags="\\Seen",  # Mark as read
-        )
-
-        assert result is True
-        # Check that append was called with the \\Seen flag
-        call_args = mock_conn.append.call_args
-        assert call_args[0][1] == "(\\Seen)"
-
-
 class TestRestoreIntegration:
     """Integration tests for restore functionality."""
 
-    def test_restore_single_folder(self, single_mock_server, monkeypatch, tmp_path):
+    def test_restore_single_folder(self, single_mock_server, tmp_path):
         """Test restoring a single folder."""
         # Create backup structure
         inbox = tmp_path / "INBOX"
@@ -356,23 +254,39 @@ Body content.
         src_data = {"INBOX": []}
         server, port = single_mock_server(src_data)
 
-        env = {
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "user",
-            "DEST_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["restore_imap_emails.py", "--src-path", str(tmp_path), "INBOX"],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
+        env = _mock_restore_env(port)
+        with temp_env(env), temp_argv(["restore_imap_emails.py", "--src-path", str(tmp_path), "INBOX"]):
+            restore_imap_emails.main()
 
-        # Run restore
-        restore_imap_emails.main()
+    def test_restore_all_folders_scans_backup(self, single_mock_server, tmp_path):
+        """End-to-end: restore all folders from a backup tree."""
+        inbox = tmp_path / "INBOX"
+        inbox.mkdir()
+        (inbox / "1_Test_Email.eml").write_text("Subject: Inbox\nMessage-ID: <inbox@test>\n\nBody")
 
-    def test_restore_single_folder_full_restore_flag(self, single_mock_server, monkeypatch, tmp_path):
+        archive = tmp_path / "Archive"
+        archive.mkdir()
+        subfolder = archive / "Sub"
+        subfolder.mkdir()
+        (subfolder / "2_Test_Email.eml").write_text("Subject: Archive\nMessage-ID: <archive@test>\n\nBody")
+
+        empty_folder = tmp_path / "Empty"
+        empty_folder.mkdir()
+
+        dest_data = {"INBOX": []}
+        server, port = single_mock_server(dest_data)
+
+        env = _mock_restore_env(port)
+        with temp_env(env), temp_argv(["restore_imap_emails.py", "--src-path", str(tmp_path)]):
+            restore_imap_emails.main()
+
+        assert "INBOX" in server.folders
+        assert "Archive/Sub" in server.folders
+        assert len(server.folders["INBOX"]) == 1
+        assert len(server.folders["Archive/Sub"]) == 1
+        assert "Empty" not in server.folders
+
+    def test_restore_single_folder_full_restore_flag(self, single_mock_server, tmp_path):
         """Smoke test: --full-restore flag is accepted."""
         inbox = tmp_path / "INBOX"
         inbox.mkdir()
@@ -390,20 +304,12 @@ Body content.
         src_data = {"INBOX": []}
         _server, port = single_mock_server(src_data)
 
-        env = {
-            "DEST_IMAP_HOST": "localhost",
-            "DEST_IMAP_USERNAME": "user",
-            "DEST_IMAP_PASSWORD": "pass",
-        }
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(
-            sys,
-            "argv",
-            ["restore_imap_emails.py", "--src-path", str(tmp_path), "--full-restore", "INBOX"],
-        )
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        restore_imap_emails.main()
+        env = _mock_restore_env(port)
+        with (
+            temp_env(env),
+            temp_argv(["restore_imap_emails.py", "--src-path", str(tmp_path), "--full-restore", "INBOX"]),
+        ):
+            restore_imap_emails.main()
 
     def test_restore_with_labels_manifest(self, tmp_path):
         """Test that labels manifest is loaded correctly."""
@@ -416,48 +322,49 @@ Body content.
         manifest_path.write_text(json.dumps(manifest_data))
 
         # Load and verify
-        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
+        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
         assert len(result) == 2
         assert result["<msg1@test.com>"] == ["INBOX", "Work"]
 
+    def test_restore_gmail_mode_fallback_folder(self, single_mock_server, tmp_path):
+        """End-to-end: Gmail mode with no labels uses fallback folder."""
+        gmail_all_mail = tmp_path / "[Gmail]" / "All Mail"
+        gmail_all_mail.mkdir(parents=True)
+        (gmail_all_mail / "1_Test.eml").write_text("Subject: X\nMessage-ID: <no-labels@test>\n\nBody")
 
-class TestEmailExistsInFolder:
-    """Tests for duplicate detection."""
+        dest_data = {"INBOX": []}
+        server, port = single_mock_server(dest_data)
 
-    def test_email_exists_true(self, monkeypatch):
-        """Test detecting existing email."""
-        mock_conn = MagicMock()
-        monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: True)
+        env = _mock_restore_env(port)
+        with temp_env(env), temp_argv(["restore_imap_emails.py", "--src-path", str(tmp_path), "--gmail-mode"]):
+            restore_imap_emails.main()
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
-        assert result is True
+        assert imap_common.FOLDER_RESTORED_UNLABELED in server.folders
+        assert len(server.folders[imap_common.FOLDER_RESTORED_UNLABELED]) == 1
 
-    def test_email_exists_false(self, monkeypatch):
-        """Test detecting non-existing email."""
-        mock_conn = MagicMock()
-        monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False)
+    def test_restore_dest_delete_removes_orphans(self, single_mock_server, tmp_path):
+        """End-to-end: --dest-delete removes messages not in local backup."""
+        inbox = tmp_path / "INBOX"
+        inbox.mkdir()
+        (inbox / "1_keep.eml").write_text("Subject: Keep\nMessage-ID: <keep@test>\n\nBody")
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
-        assert result is False
+        dest_data = {
+            "INBOX": [
+                b"Subject: Keep\r\nMessage-ID: <keep@test>\r\n\r\nBody",
+                b"Subject: Orphan\r\nMessage-ID: <orphan@test>\r\n\r\nBody",
+            ]
+        }
+        server, port = single_mock_server(dest_data)
 
-    def test_email_exists_no_message_id(self, monkeypatch):
-        """Test with no message ID."""
-        mock_conn = MagicMock()
+        env = _mock_restore_env(port)
+        with (
+            temp_env(env),
+            temp_argv(["restore_imap_emails.py", "--src-path", str(tmp_path), "--dest-delete", "INBOX"]),
+        ):
+            restore_imap_emails.main()
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, None)
-        assert result is False
-
-    def test_email_exists_exception(self, monkeypatch):
-        """Test handling exception."""
-        mock_conn = MagicMock()
-
-        monkeypatch.setattr(
-            "imap_common.message_exists_in_folder",
-            lambda *args: (_ for _ in ()).throw(Exception("Error")),
-        )
-
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
-        assert result is False
+        assert len(server.folders["INBOX"]) == 1
+        assert b"Message-ID: <keep@test>" in server.folders["INBOX"][0]["content"]
 
 
 class TestRestoreProgressCache:
@@ -500,6 +407,60 @@ class TestRestoreProgressCache:
         assert "<b@test>" in ids
 
 
+class TestBackupFolderDiscovery:
+    """Tests for backup folder discovery helpers."""
+
+    def test_get_backup_folders_skips_unreadable(self, tmp_path):
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1.eml").write_bytes(b"Subject: Inbox\r\n\r\nBody")
+
+        parent_path = tmp_path / "Parent"
+        parent_path.mkdir()
+
+        unreadable_path = parent_path / "Unreadable"
+        unreadable_path.mkdir()
+        (unreadable_path / "1.eml").write_bytes(b"Subject: Hidden\r\n\r\nBody")
+        os.chmod(unreadable_path, 0)
+
+        try:
+            folders = imap_common.get_backup_folders(str(tmp_path))
+        finally:
+            os.chmod(unreadable_path, 0o700)
+
+        folder_names = {name for name, _path in folders}
+        assert "INBOX" in folder_names
+        assert "Unreadable" not in folder_names
+
+    def test_extract_message_id_from_eml_missing_file(self, tmp_path):
+        missing_path = tmp_path / "missing.eml"
+        assert imap_common.extract_message_id_from_eml(str(missing_path)) is None
+
+    def test_extract_message_id_from_eml_success(self, tmp_path):
+        eml_path = tmp_path / "message.eml"
+        eml_path.write_text("Message-ID: <ok@test>\r\nSubject: Hi\r\n\r\nBody")
+
+        assert imap_common.extract_message_id_from_eml(str(eml_path)) == "<ok@test>"
+
+
+class TestTrashFolderDetection:
+    """Tests for trash folder detection with string LIST entries."""
+
+    def test_detect_trash_folder_with_string_entries(self):
+        class FakeConn:
+            def list(self):
+                return (
+                    "OK",
+                    [
+                        '(\\HasNoChildren) "/" "INBOX"',
+                        '(\\HasNoChildren \\Trash) "/" "Trash"',
+                    ],
+                )
+
+        result = imap_common.detect_trash_folder(FakeConn())
+        assert result == "Trash"
+
+
 class TestGetLabelsFromManifest:
     """Tests for get_labels_from_manifest function."""
 
@@ -532,107 +493,48 @@ class TestGetLabelsFromManifest:
         assert result == []
 
 
-class TestLabelToFolder:
-    """Tests for label_to_folder function."""
+class TestRestoreE2EHelpers:
+    """End-to-end tests for restore helper functions using the mock IMAP server."""
 
-    def test_inbox_label(self):
-        """Test INBOX label conversion."""
-        result = restore_imap_emails.label_to_folder("INBOX")
-        assert result == "INBOX"
+    def test_upload_email_success(self, single_mock_server):
+        dest_data = {"INBOX": []}
+        server, port = single_mock_server(dest_data)
 
-    def test_sent_mail_label(self):
-        """Test Sent Mail label conversion."""
-        result = restore_imap_emails.label_to_folder("Sent Mail")
-        assert result == "[Gmail]/Sent Mail"
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
 
-    def test_starred_label(self):
-        """Test Starred label conversion."""
-        result = restore_imap_emails.label_to_folder("Starred")
-        assert result == "[Gmail]/Starred"
-
-    def test_drafts_label(self):
-        """Test Drafts label conversion."""
-        result = restore_imap_emails.label_to_folder("Drafts")
-        assert result == "[Gmail]/Drafts"
-
-    def test_important_label(self):
-        """Test Important label conversion."""
-        result = restore_imap_emails.label_to_folder("Important")
-        assert result == "[Gmail]/Important"
-
-    def test_custom_label(self):
-        """Test custom user label (no conversion)."""
-        result = restore_imap_emails.label_to_folder("Work")
-        assert result == "Work"
-
-    def test_nested_label(self):
-        """Test nested user label (no conversion)."""
-        result = restore_imap_emails.label_to_folder("Projects/2024")
-        assert result == "Projects/2024"
-
-
-class TestGmailModeDraftsFallbackRegression:
-    def test_gmail_mode_no_labels_does_not_upload_to_drafts(self, monkeypatch):
-        """Regression: messages with no usable labels must not be uploaded to Gmail Drafts."""
-
-        captured = {}
-
-        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
-
-        def fake_parse_eml_file(_path):
-            return (
-                "<no-labels@test>",
-                '"01-Jan-2024 00:00:00 +0000"',
-                b"Subject: X\r\nMessage-ID: <no-labels@test>\r\n\r\nBody",
-                "X",
-            )
-
-        monkeypatch.setattr(restore_imap_emails, "get_thread_connection", lambda _conf: MagicMock())
-        monkeypatch.setattr(restore_imap_emails, "upload_email", fake_upload_email)
-        monkeypatch.setattr(restore_imap_emails, "parse_eml_file", fake_parse_eml_file)
-
-        restore_imap_emails.process_restore_batch(
-            eml_files=[("/does/not/matter.eml", "x.eml")],
-            folder_name="__GMAIL_MODE__",
-            dest_conf=("host", "user", "pass"),
-            manifest={},
-            apply_labels=True,
-            apply_flags=False,
+        result = restore_imap_emails.upload_email(
+            conn,
+            "INBOX",
+            b"Subject: Upload\r\nMessage-ID: <up@test>\r\n\r\nBody",
+            '"15-Jan-2024 10:30:00 +0000"',
         )
 
-        assert captured["folder_name"] == imap_common.FOLDER_RESTORED_UNLABELED
-        assert captured["folder_name"] != "[Gmail]/Drafts"
+        assert result == restore_imap_emails.UploadResult.SUCCESS
+        assert len(server.folders["INBOX"]) == 1
+        conn.logout()
 
+    def test_sync_flags_on_existing(self, single_mock_server):
+        dest_data = {
+            "INBOX": [{"uid": 1, "flags": set(), "content": b"Subject: Flag\r\nMessage-ID: <flag@test>\r\n\r\nBody"}]
+        }
+        _server, port = single_mock_server(dest_data)
 
-class TestSyncFlagsOnExisting:
-    """Tests for sync_flags_on_existing function."""
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
 
-    def test_sync_flags_adds_missing(self):
-        """Test that missing flags are added to existing email."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.search.return_value = ("OK", [b"1"])
-        mock_conn.fetch.return_value = ("OK", [(b"1 (FLAGS ())", b"")])
-        mock_conn.store.return_value = ("OK", None)
+        imap_common.sync_flags_on_existing(conn, "INBOX", "<flag@test>", "\\Seen \\Flagged", 1000)
 
-        # Should not raise
-        restore_imap_emails.sync_flags_on_existing(mock_conn, "INBOX", "<test@test.com>", "\\Seen \\Flagged", 1000)
-
-        # Verify store was called with flags
-        mock_conn.store.assert_called()
-
-    def test_sync_flags_no_message_found(self):
-        """Test when message is not found."""
-        mock_conn = MagicMock()
-        mock_conn.select.return_value = ("OK", [b"1"])
-        mock_conn.search.return_value = ("OK", [b""])  # No match
-
-        # Should not raise or call store
-        restore_imap_emails.sync_flags_on_existing(mock_conn, "INBOX", "<test@test.com>", "\\Seen", 1000)
-
-        mock_conn.store.assert_not_called()
+        conn.select('"INBOX"')
+        resp, data = conn.search(None, 'HEADER Message-ID "<flag@test>"')
+        assert resp == "OK"
+        msg_num = data[0].split()[0]
+        resp, flag_data = conn.fetch(msg_num, "(FLAGS)")
+        assert resp == "OK"
+        flag_text = str(flag_data[0])
+        assert "\\Seen" in flag_text
+        assert "\\Flagged" in flag_text
+        conn.logout()
 
 
 class TestDestDeleteRestoreArgument:
@@ -683,7 +585,7 @@ class TestDestDeleteRestoreFunctionality:
         # Local backup only has one email
         local_msg_ids = {"<keep@test>"}
 
-        deleted = restore_imap_emails.delete_orphan_emails_from_dest(conn, "INBOX", local_msg_ids)
+        deleted = imap_common.delete_orphan_emails(conn, "INBOX", local_msg_ids)
 
         assert deleted == 2
         # Verify only 1 email remains
@@ -711,7 +613,7 @@ class TestDestDeleteRestoreFunctionality:
         # Empty local backup
         local_msg_ids = set()
 
-        deleted = restore_imap_emails.delete_orphan_emails_from_dest(conn, "INBOX", local_msg_ids)
+        deleted = imap_common.delete_orphan_emails(conn, "INBOX", local_msg_ids)
 
         assert deleted == 2
         assert len(server.folders["INBOX"]) == 0
@@ -737,7 +639,7 @@ class TestDestDeleteRestoreFunctionality:
         # Local has both emails
         local_msg_ids = {"<e1@test>", "<e2@test>"}
 
-        deleted = restore_imap_emails.delete_orphan_emails_from_dest(conn, "INBOX", local_msg_ids)
+        deleted = imap_common.delete_orphan_emails(conn, "INBOX", local_msg_ids)
 
         assert deleted == 0
         assert len(server.folders["INBOX"]) == 2
@@ -784,9 +686,7 @@ class TestDestDeleteRestoreFunctionality:
         assert len(msg_ids) == 1
         assert "<valid@test>" in msg_ids
 
-    def test_dest_delete_enabled_via_env_var_main_deletes_all_in_folder(
-        self, single_mock_server, monkeypatch, tmp_path
-    ):
+    def test_dest_delete_enabled_via_env_var_main_deletes_all_in_folder(self, single_mock_server, tmp_path):
         """End-to-end: DEST_DELETE=true triggers deletion when local folder has no .eml files."""
         dest_data = {
             "INBOX": [
@@ -800,17 +700,74 @@ class TestDestDeleteRestoreFunctionality:
         backup_root = tmp_path / "backup"
         (backup_root / "INBOX").mkdir(parents=True)
 
-        monkeypatch.setenv("BACKUP_LOCAL_PATH", str(backup_root))
-        monkeypatch.setenv("DEST_IMAP_HOST", "localhost")
-        monkeypatch.setenv("DEST_IMAP_USERNAME", "user")
-        monkeypatch.setenv("DEST_IMAP_PASSWORD", "pass")
-        monkeypatch.setenv("MAX_WORKERS", "1")
-        monkeypatch.setenv("DEST_DELETE", "true")
-
-        # Force folder-mode so empty folders still get processed.
-        monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "INBOX"])
-        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
-
-        restore_imap_emails.main()
+        env = _mock_restore_env(port)
+        env.update(
+            {
+                "BACKUP_LOCAL_PATH": str(backup_root),
+                "DEST_DELETE": "true",
+            }
+        )
+        with temp_env(env), temp_argv(["restore_imap_emails.py", "INBOX"]):
+            restore_imap_emails.main()
 
         assert len(server.folders["INBOX"]) == 0
+
+
+class TestAppendEmailReturnValueChecking:
+    """Tests that verify append_email return values are checked before recording progress."""
+
+    def test_record_progress_not_called_on_append_failure(self, tmp_path):
+        """Test that record_progress is not called when append_email fails during label application."""
+        from unittest.mock import Mock, patch
+
+        # Create test email file
+        backup_root = tmp_path / "backup"
+        inbox = backup_root / "INBOX"
+        inbox.mkdir(parents=True)
+
+        eml_content = b"""From: sender@test.com
+To: recipient@test.com
+Subject: Test Email
+Message-ID: <test123@test.com>
+Date: Mon, 15 Jan 2024 10:30:00 +0000
+
+Body content.
+"""
+        (inbox / "1_Test_Email.eml").write_bytes(eml_content)
+
+        # Create labels manifest with multiple labels
+        manifest = {"<test123@test.com>": {"labels": ["INBOX", "Work"], "flags": []}}
+        manifest_path = backup_root / "labels_manifest.json"
+        manifest_path.write_text(json.dumps(manifest))
+
+        # Mock IMAP connection that fails on the second append (for label)
+        mock_conn = Mock()
+        mock_conn.select.return_value = ("OK", [b"0"])
+        mock_conn.search.return_value = ("OK", [b""])  # No duplicates
+
+        # First append succeeds (INBOX), second append fails (Work label)
+        mock_conn.append.side_effect = [
+            ("OK", []),  # INBOX upload succeeds
+            ("NO", []),  # Work label append fails
+        ]
+
+        # Track calls to record_progress
+        with (
+            patch("utils.restore_cache.record_progress") as mock_record_progress,
+            patch("utils.imap_common.get_imap_connection", return_value=mock_conn),
+        ):
+            env = {
+                "DEST_IMAP_HOST": "localhost",
+                "DEST_IMAP_USERNAME": "user",
+                "DEST_IMAP_PASSWORD": "pass",
+                "MAX_WORKERS": "1",
+            }
+            with temp_env(env), temp_argv(["restore_imap_emails.py", "--src-path", str(backup_root), "INBOX"]):
+                restore_imap_emails.main()
+
+            # Verify record_progress was called only once (for successful INBOX upload)
+            # and NOT called for the failed Work label append
+            assert mock_record_progress.call_count == 1
+            # Verify it was called for INBOX only
+            call_args = mock_record_progress.call_args
+            assert call_args[1]["folder_name"] == "INBOX"

test/test_imap_retry.py

@@ -0,0 +1,188 @@
+"""
+Tests for imap_retry.py
+
+Tests cover:
+- Transient error detection
+- ConnectionProxy transparent proxying
+- Retry with exponential backoff on transient errors
+- Pass-through for non-retryable methods
+- Pass-through for non-transient errors
+"""
+
+import imaplib
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tools")))
+
+from mock_imap_server import start_server_thread as start_mock_server
+from utils import imap_retry
+
+
+class TestConnectionProxy:
+    @pytest.fixture(scope="class")
+    def imap_server_info(self):
+        folders = {"INBOX": []}
+        server, port = start_mock_server(folders)
+        yield server, port
+        server.shutdown()
+        server.server_close()
+
+    @pytest.fixture
+    def imap_conn(self, imap_server_info):
+        server, port = imap_server_info
+        client = imaplib.IMAP4("127.0.0.1", port)
+        client.login("user", "pass")
+        yield client
+        try:
+            client.logout()
+        except:
+            pass
+
+    @pytest.fixture
+    def captured_logs(self):
+        logs = []
+
+        def log_fn(msg):
+            logs.append(msg)
+
+        return logs, log_fn
+
+    @pytest.fixture
+    def proxy(self, imap_conn, captured_logs):
+        # Short wait to speed up tests
+        logs, log_fn = captured_logs
+        return imap_retry.ConnectionProxy(imap_conn, max_retries=3, initial_wait=0.01, log_fn=log_fn)
+
+    def test_transient_retry_logic_unavailable(self, proxy, captured_logs):
+        """Test detection of 'NO [UNAVAILABLE]' error pattern."""
+        try:
+            logs, _ = captured_logs
+            # Server will echo "NO [UNAVAILABLE] Server Busy" and fail continuously
+            typ, data = proxy.select('"NO [UNAVAILABLE] Server Busy"')
+            assert typ == "NO"
+            # Should have retried max_retries-1 times before returning last error
+            assert len(logs) == 2  # Retries on attempt 1 and 2, then fail on 3
+            assert "attempt 1/3" in logs[0]
+            assert "attempt 2/3" in logs[1]
+        except Exception:
+            raise
+
+    def test_transient_retry_logic_server_busy(self, proxy, captured_logs):
+        """Test detection of 'NO Server Busy' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO Server Busy"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_transient_retry_logic_try_again(self, proxy, captured_logs):
+        """Test detection of 'try again' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO try again later"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_transient_retry_logic_throttled(self, proxy, captured_logs):
+        """Test detection of 'THROTTLED' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO [THROTTLED]"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_non_transient_no_retry_empty(self, proxy, captured_logs):
+        """Test mismatch on empty/irrelevant data."""
+        logs, _ = captured_logs
+        # EMPTY arg causes mock to return OK with 0 items but generally here we want an error condition
+        # Mock server "EMPTY" folder returns OK.
+        # We need an error that is NOT transient.
+        # "NO [UNKNOWN]"
+        typ, data = proxy.select('"NO [UNKNOWN]"')
+        assert typ == "NO"
+        assert len(logs) == 0
+
+    def test_ok_response_returns_immediately(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.select('"INBOX"')
+        assert typ == "OK"
+        assert len(logs) == 0
+
+    def test_non_transient_error_not_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO [AUTHENTICATIONFAILED]"')
+        assert typ == "NO"
+        assert b"AUTHENTICATIONFAILED" in data[0]
+        assert len(logs) == 0
+
+    def test_transient_error_retried_then_succeeds(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        # Retry 2 times (fail count 2, so 3rd succeeds)
+        # RETRY_2 triggers 2 failures then OK.
+        typ, data = proxy.select('"RETRY_2"')
+        assert typ == "OK"
+        assert len(logs) == 2
+        assert "attempt 1/3" in logs[0]
+        assert "attempt 2/3" in logs[1]
+
+    def test_exponential_backoff_logic(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        typ, data = proxy.fetch("1", "(BODY RETRY_2)")
+        assert typ == "OK"
+        assert len(logs) == 2
+        # Verify message content for backoff?
+        # approximate wait time is hard to verify without mocking sleep, but we verify calls were made.
+
+    def test_max_retries_exhausted_returns_last_error(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        # Fail 5 times. Max retries is 3. Result should be error.
+        typ, data = proxy.store("1", "+FLAGS", "(\\Seen RETRY_5)")
+        assert typ == "NO"
+        assert b"UNAVAILABLE" in data[0]
+        assert len(logs) == 2
+
+    def test_non_retryable_method_passes_through(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        res = proxy.capability()
+        assert res[0] == "OK"
+        assert len(logs) == 0
+
+    def test_non_callable_attribute_passes_through(self, proxy):
+        assert proxy.state == "AUTH"
+
+    def test_non_tuple_return_passes_through(self):
+        class DummyConn:
+            def noop(self):
+                return "unexpected"
+
+        d = DummyConn()
+        p = imap_retry.ConnectionProxy(d)
+        assert p.noop() == "unexpected"
+
+    def test_uid_method_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        typ, data = proxy.uid("SEARCH", "RETRY_1")
+        assert typ == "OK"
+        assert len(logs) == 1
+
+    def test_append_method_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.append('"RETRY_1"', None, None, b"data")
+        assert typ == "OK"
+        assert len(logs) == 1
+
+    def test_max_retries_zero_raises(self):
+        with pytest.raises(ValueError, match="max_retries must be >= 1"):
+            imap_retry.ConnectionProxy(None, max_retries=0)
+
+    def test_max_retries_negative_raises(self):
+        with pytest.raises(ValueError, match="max_retries must be >= 1"):
+            imap_retry.ConnectionProxy(None, max_retries=-1)
+
+    def test_initial_wait_negative_raises(self):
+        with pytest.raises(ValueError, match="initial_wait must be >= 0"):
+            imap_retry.ConnectionProxy(None, initial_wait=-1)

test/test_migrate_resume.py

@@ -0,0 +1,165 @@
+"""Tests for migration resumption functionality using cache."""
+
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_migrate as migrate_imap_emails
+from conftest import temp_argv, temp_env
+from utils import restore_cache
+
+
+def _run_migrate(cache_dir, src_port, dest_port):
+    env = {
+        "SRC_IMAP_HOST": f"imap://localhost:{src_port}",
+        "SRC_IMAP_USERNAME": "src",
+        "SRC_IMAP_PASSWORD": "p",
+        "DEST_IMAP_HOST": f"imap://localhost:{dest_port}",
+        "DEST_IMAP_USERNAME": "dest",
+        "DEST_IMAP_PASSWORD": "p",
+        "MAX_WORKERS": "1",
+    }
+
+    argv = [
+        "migrate_imap_emails.py",
+        "--src-host",
+        f"imap://localhost:{src_port}",
+        "--src-user",
+        "src",
+        "--src-pass",
+        "p",
+        "--dest-host",
+        f"imap://localhost:{dest_port}",
+        "--dest-user",
+        "dest",
+        "--dest-pass",
+        "p",
+        "--migrate-cache",
+        str(cache_dir),
+        "--workers",
+        "1",
+    ]
+
+    with temp_env(env), temp_argv(argv):
+        migrate_imap_emails.main()
+
+
+class TestMigrationResumption:
+    """Tests that migration resumes from the last known UID using source state tracking."""
+
+    def test_migrate_resumes_using_last_uid(self, mock_server_factory, tmp_path):
+        # Setup source with 3 messages
+        msg1 = b"Subject: Msg1\r\nMessage-ID: <msg1@test>\r\n\r\nBody1"
+        msg2 = b"Subject: Msg2\r\nMessage-ID: <msg2@test>\r\n\r\nBody2"
+        msg3 = b"Subject: Msg3\r\nMessage-ID: <msg3@test>\r\n\r\nBody3"
+        src_data = {"INBOX": [msg1, msg2, msg3]}  # UIDs will be 1, 2, 3
+        dest_data = {"INBOX": []}
+
+        _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        cache_dir = tmp_path / "cache"
+        cache_dir.mkdir()
+
+        # Pre-seed cache to simulate a prior run that finished up to UID 2
+        # Mock server uses UIDVALIDITY=1 by default
+        dest_host = f"imap://localhost:{p2}"
+        dest_user = "dest"
+        cache_path = restore_cache.get_dest_index_cache_path(str(cache_dir), dest_host, dest_user)
+
+        cache_data = {
+            "version": 1,
+            "dest": {"host": dest_host, "user": dest_user},
+            "folders": {
+                "INBOX": {
+                    # Simulate that we've already seen msg1 and msg2
+                    "message_ids": ["<msg1@test>", "<msg2@test>"],
+                    "source_state": {"uid_validity": 1, "last_uid": 2},
+                }
+            },
+            "_meta": {},
+        }
+
+        # Update cache data for the robust test
+        # We delete msg1 from cache ID list but keep last_uid=2.
+        # If it resumes from 2, it won't see msg1 (UID 1).
+        cache_data["folders"]["INBOX"]["message_ids"] = ["<msg2@test>"]  # Removed msg1
+
+        with open(cache_path, "w") as f:
+            json.dump(cache_data, f)
+
+        # Run without patch - relies on real localhost tcp connection
+        _run_migrate(cache_dir, p1, p2)
+
+        msgs_in_dest = dest_server.folders["INBOX"]
+        assert len(msgs_in_dest) == 1
+        assert b"Msg3" in msgs_in_dest[0]["content"]
+
+        # Check that cache was updated with new last_uid
+        with open(cache_path) as f:
+            new_cache = json.load(f)
+
+        src_state = new_cache["folders"]["INBOX"]["source_state"]
+        assert src_state["last_uid"] == 3
+        # Should also have added msg3 to message_ids
+        assert "<msg3@test>" in new_cache["folders"]["INBOX"]["message_ids"]
+
+    def test_migrate_ignores_resume_on_uidvalidity_mismatch(self, mock_server_factory, tmp_path):
+        # Setup source with 2 messages
+        msg1 = b"Subject: Msg1\r\nMessage-ID: <msg1@test>\r\n\r\nBody1"
+        msg2 = b"Subject: Msg2\r\nMessage-ID: <msg2@test>\r\n\r\nBody2"
+        src_data = {"INBOX": [msg1, msg2]}  # UIDs 1, 2
+        dest_data = {"INBOX": []}
+
+        _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        cache_dir = tmp_path / "cache"
+        cache_dir.mkdir()
+
+        dest_host = f"imap://localhost:{p2}"
+        dest_user = "dest"
+        cache_path = restore_cache.get_dest_index_cache_path(str(cache_dir), dest_host, dest_user)
+
+        # Pre-seed cache with mismatching UIDVALIDITY (999 vs 1)
+        # But claim we processed everything (last_uid=2)
+        # Also remove msg1 from cache to detect if it gets re-processed
+        cache_data = {
+            "version": 1,
+            "dest": {"host": dest_host, "user": dest_user},
+            "folders": {
+                "INBOX": {
+                    "message_ids": ["<msg2@test>"],
+                    "source_state": {
+                        "uid_validity": 999,  # Mismatch
+                        "last_uid": 2,
+                    },
+                }
+            },
+            "_meta": {},
+        }
+
+        with open(cache_path, "w") as f:
+            json.dump(cache_data, f)
+
+        # Run without patch - relies on real localhost tcp connection
+        _run_migrate(cache_dir, p1, p2)
+
+        # Since UIDVALIDITY mismatched, it should ignore last_uid=2 and rescan from start (UID 1 and 2).
+        # UID 1 is NOT in cache -> should be copied.
+        # UID 2 IS in cache -> should be skipped.
+
+        msgs_in_dest = dest_server.folders["INBOX"]
+        assert len(msgs_in_dest) == 1
+        assert b"Msg1" in msgs_in_dest[0]["content"]
+
+        # Verify it updated the cache to the NEW UIDVALIDITY (1) and last_uid (2)
+        with open(cache_path) as f:
+            new_cache = json.load(f)
+
+        src_state = new_cache["folders"]["INBOX"]["source_state"]
+        assert src_state["uid_validity"] == 1
+        # last_uid is 1 because UID 2 was skipped as a duplicate (pre-filtered),
+        # so the batch processor only saw UID 1.
+        # This is acceptable (conservative) behavior.
+        assert src_state["last_uid"] == 1

test/test_migrate_with_cache.py

@@ -0,0 +1,248 @@
+"""End-to-end tests for migrate_imap_emails.py cache behavior."""
+
+import imaplib
+import json
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_migrate as migrate_imap_emails
+from conftest import make_mock_connection, temp_argv, temp_env
+from core import imap_session
+from utils import imap_common, restore_cache
+
+
+def _run_migrate(cache_dir, src_port, dest_port, full_migrate=False, extra_env=None):
+    env = {
+        "SRC_IMAP_HOST": f"imap://localhost:{src_port}",
+        "SRC_IMAP_USERNAME": "src",
+        "SRC_IMAP_PASSWORD": "p",
+        "DEST_IMAP_HOST": f"imap://localhost:{dest_port}",
+        "DEST_IMAP_USERNAME": "dest",
+        "DEST_IMAP_PASSWORD": "p",
+        "MAX_WORKERS": "1",
+    }
+    if extra_env:
+        env.update(extra_env)
+
+    argv = [
+        "migrate_imap_emails.py",
+        "--src-host",
+        f"imap://localhost:{src_port}",
+        "--src-user",
+        "src",
+        "--src-pass",
+        "p",
+        "--dest-host",
+        f"imap://localhost:{dest_port}",
+        "--dest-user",
+        "dest",
+        "--dest-pass",
+        "p",
+        "--migrate-cache",
+        str(cache_dir),
+        "--workers",
+        "1",
+    ]
+    if full_migrate:
+        argv.append("--full-migrate")
+
+    with temp_env(env), temp_argv(argv):
+        migrate_imap_emails.main()
+
+
+@pytest.mark.usefixtures("mock_server_factory")
+class TestMigrationCache:
+    """End-to-end tests for incremental migration using local cache."""
+
+    def test_migrate_skips_cached_items(self, mock_server_factory, tmp_path):
+        msg_id = "<cached@test>"
+        msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": []}
+
+        _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        cache_dir = tmp_path / "cache"
+        cache_dir.mkdir()
+
+        # First run populates cache and copies message.
+        _run_migrate(cache_dir, p1, p2)
+        assert len(dest_server.folders["INBOX"]) == 1
+
+        # Second run against a fresh destination should skip based on cache.
+        dest_server.folders["INBOX"] = []
+        _run_migrate(cache_dir, p1, p2)
+        assert len(dest_server.folders["INBOX"]) == 0
+
+    def test_migrate_writes_to_cache(self, mock_server_factory, tmp_path):
+        msg_id = "<new@test>"
+        msg = f"Subject: New\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": []}
+
+        _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        cache_dir = tmp_path / "cache"
+        cache_dir.mkdir()
+
+        _run_migrate(cache_dir, p1, p2)
+
+        assert len(dest_server.folders["INBOX"]) == 1
+
+        cache_path = restore_cache.get_dest_index_cache_path(str(cache_dir), f"imap://localhost:{p2}", "dest")
+        assert os.path.exists(cache_path)
+
+        with open(cache_path, encoding="utf-8") as f:
+            cache_data = json.load(f)
+
+        msg_ids = set(cache_data.get("folders", {}).get("INBOX", {}).get("message_ids", []))
+        assert msg_id in msg_ids
+
+    def test_full_migrate_ignores_cache(self, mock_server_factory, tmp_path):
+        msg_id = "<cached@test>"
+        msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+
+        cache_dir = tmp_path / "cache"
+        cache_dir.mkdir()
+
+        # Populate cache with an initial run.
+        _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, {"INBOX": []})
+        _run_migrate(cache_dir, p1, p2)
+
+        # Fresh destination should still copy when --full-migrate is set.
+        _dest_server.folders["INBOX"] = []
+        _run_migrate(cache_dir, p1, p2, full_migrate=True)
+
+        assert len(_dest_server.folders["INBOX"]) == 1
+
+    def test_cached_skip_with_preserve_flags(self, mock_server_factory, tmp_path):
+        msg_id = "<cached-preserve@test>"
+        msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+
+        cache_dir = tmp_path / "cache"
+        cache_dir.mkdir()
+
+        # Populate cache with initial run.
+        _src_server, _dest_server, p1, p2 = mock_server_factory(src_data, {"INBOX": []})
+        _run_migrate(cache_dir, p1, p2)
+
+        # Preserve flags disables pre-filtering, so cache skip happens per message.
+        _dest_server.folders["INBOX"] = []
+        _run_migrate(cache_dir, p1, p2, extra_env={"PRESERVE_FLAGS": "true"})
+
+        assert len(_dest_server.folders["INBOX"]) == 0
+
+    def test_load_progress_cache_warns_on_unusable_root(self, tmp_path):
+        cache_file = tmp_path / "cachefile"
+        cache_file.write_text("not a directory")
+
+        messages = []
+        _cache_path, _cache_data, _cache_lock = imap_common.load_progress_cache(
+            str(cache_file),
+            "host",
+            "user",
+            log_fn=messages.append,
+        )
+
+        assert any("unable to create cache directory" in msg for msg in messages)
+
+    def test_migrate_folder_logs_cache_load_failure(self, mock_server_factory, tmp_path, capsys):
+        msg_id = "<cache-fail@test>"
+        msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": []}
+
+        _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        src = imaplib.IMAP4("localhost", p1)
+        dest = imaplib.IMAP4("localhost", p2)
+        src.login("src", "p")
+        dest.login("dest", "p")
+
+        with (
+            patch("utils.imap_common.get_imap_connection", make_mock_connection(p1, p2)),
+            patch.object(
+                imap_common,
+                "load_progress_cache",
+                side_effect=RuntimeError("boom"),
+            ),
+            patch.object(
+                imap_session, "get_thread_connection", lambda _store, key, _conf: src if key == "src" else dest
+            ),
+        ):
+            migrate_imap_emails.MAX_WORKERS = 1
+            migrate_imap_emails.BATCH_SIZE = 1
+
+            migrate_imap_emails.migrate_folder(
+                src,
+                dest,
+                "INBOX",
+                False,
+                {"host": "localhost", "user": "src", "password": "p"},
+                {"host": "localhost", "user": "dest", "password": "p"},
+                progress_cache_path=str(tmp_path / "cache"),
+                progress_cache_file=None,
+                progress_cache_data=None,
+                progress_cache_lock=None,
+            )
+
+        captured = capsys.readouterr()
+        assert "Warning: Failed to load cache" in captured.out
+        assert len(dest_server.folders["INBOX"]) == 1
+
+        src.logout()
+        dest.logout()
+
+    def test_migrate_folder_logs_cache_read_failure(self, mock_server_factory, tmp_path, capsys):
+        msg_id = "<cache-read-fail@test>"
+        msg = f"Subject: Cached\r\nMessage-ID: {msg_id}\r\n\r\nBody".encode()
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": []}
+
+        _src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        src = imaplib.IMAP4("localhost", p1)
+        dest = imaplib.IMAP4("localhost", p2)
+        src.login("src", "p")
+        dest.login("dest", "p")
+
+        with (
+            patch("utils.imap_common.get_imap_connection", make_mock_connection(p1, p2)),
+            patch.object(
+                restore_cache,
+                "get_cached_message_ids",
+                side_effect=RuntimeError("read fail"),
+            ),
+            patch.object(
+                imap_session, "get_thread_connection", lambda _store, key, _conf: src if key == "src" else dest
+            ),
+        ):
+            migrate_imap_emails.MAX_WORKERS = 1
+            migrate_imap_emails.BATCH_SIZE = 1
+
+            migrate_imap_emails.migrate_folder(
+                src,
+                dest,
+                "INBOX",
+                False,
+                {"host": "localhost", "user": "src", "password": "p"},
+                {"host": "localhost", "user": "dest", "password": "p"},
+                progress_cache_path=str(tmp_path / "cache"),
+                progress_cache_file=None,
+                progress_cache_data=None,
+                progress_cache_lock=None,
+            )
+
+        captured = capsys.readouterr()
+        assert "Warning: Failed to read cache for folder 'INBOX'" in captured.out
+        assert len(dest_server.folders["INBOX"]) == 1
+
+        src.logout()
+        dest.logout()

test/utils/test_imap_common.py

@@ -14,41 +14,39 @@ Tests cover:
 
 import os
 import sys
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+import pytest
 
-import imap_common
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+
+from conftest import temp_env
+from utils import imap_common
 
 
 class TestVerifyEnvVars:
     """Tests for verify_env_vars function."""
 
-    def test_all_vars_present(self, monkeypatch):
+    def test_all_vars_present(self):
         """Test returns True when all variables are set."""
-        monkeypatch.setenv("VAR1", "value1")
-        monkeypatch.setenv("VAR2", "value2")
+        with temp_env({"VAR1": "value1", "VAR2": "value2"}):
+            result = imap_common.verify_env_vars(["VAR1", "VAR2"])
+            assert result is True
 
-        result = imap_common.verify_env_vars(["VAR1", "VAR2"])
-        assert result is True
-
-    def test_missing_vars(self, monkeypatch, capsys):
+    def test_missing_vars(self, capsys):
         """Test returns False and prints error when variables are missing."""
-        monkeypatch.delenv("MISSING_VAR", raising=False)
+        with temp_env({}):
+            result = imap_common.verify_env_vars(["MISSING_VAR"])
+            assert result is False
 
-        result = imap_common.verify_env_vars(["MISSING_VAR"])
-        assert result is False
+            captured = capsys.readouterr()
+            assert "MISSING_VAR" in captured.err
 
-        captured = capsys.readouterr()
-        assert "MISSING_VAR" in captured.err
-
-    def test_partial_vars_present(self, monkeypatch, capsys):
+    def test_partial_vars_present(self, capsys):
         """Test with some variables present and some missing."""
-        monkeypatch.setenv("PRESENT", "value")
-        monkeypatch.delenv("MISSING", raising=False)
-
-        result = imap_common.verify_env_vars(["PRESENT", "MISSING"])
-        assert result is False
+        with temp_env({"PRESENT": "value"}):
+            result = imap_common.verify_env_vars(["PRESENT", "MISSING"])
+            assert result is False
 
 
 class TestGetImapConnection:
@@ -74,6 +72,60 @@ class TestGetImapConnection:
         captured = capsys.readouterr()
         assert "Connection error" in captured.out or "Error" in captured.out
 
+    @patch("imaplib.IMAP4")
+    @patch("imaplib.IMAP4_SSL")
+    def test_scheme_parsing(self, mock_ssl, mock_imap):
+        """Test scheme parsing determines SSL usage."""
+        # Setup mocks
+        mock_conn = Mock()
+        mock_ssl.return_value = mock_conn
+        mock_imap.return_value = mock_conn
+
+        # Test SSL schemes
+        for scheme in ["imaps", "imap+ssl", "imapssl", "ssl"]:
+            host = f"{scheme}://mail.example.com"
+            # Note: mock_conn.login is called, so we can mock it to succeed or ignore it
+            imap_common.get_imap_connection(host, "user", "pass")
+            mock_ssl.assert_called_with("mail.example.com")
+            mock_imap.assert_not_called()
+            mock_ssl.reset_mock()
+            mock_imap.reset_mock()
+
+        # Test non-SSL schemes
+        for scheme in ["imap", "tcp"]:
+            host = f"{scheme}://mail.example.com"
+            imap_common.get_imap_connection(host, "user", "pass")
+            mock_imap.assert_called_with("mail.example.com")
+            mock_ssl.assert_not_called()
+            mock_ssl.reset_mock()
+            mock_imap.reset_mock()
+
+        # Test with port
+        host = "imap://mail.example.com:143"
+        imap_common.get_imap_connection(host, "user", "pass")
+        mock_imap.assert_called_with("mail.example.com", 143)
+
+    def test_invalid_schemes(self, capsys):
+        """Test invalid schemes raise ValueError caught by exception handler."""
+        imap_common.get_imap_connection("http://mail.example.com", "user", "pass")
+        captured = capsys.readouterr()
+        assert "Unsupported IMAP scheme: http" in captured.out
+
+    def test_invalid_host_url(self, capsys):
+        """Test invalid host URL parsing."""
+        # No hostname
+        imap_common.get_imap_connection("imaps://", "user", "pass")
+        captured = capsys.readouterr()
+        assert "Invalid IMAP host" in captured.out
+
+        # No scheme (urllib fails to parse scheme if missing :// or just host)
+        # But our code checks "://" in host. If not, it assumes basic host.
+        # So we test a case where "://" is present but scheme is empty?
+        # "://hostname" parses scheme as empty string.
+        imap_common.get_imap_connection("://mail.example.com", "user", "pass")
+        captured = capsys.readouterr()
+        assert "Invalid IMAP host" in captured.out
+
 
 class TestNormalizeFolderName:
     """Tests for normalize_folder_name function."""
@@ -268,48 +320,6 @@ class TestDetectTrashFolder:
         assert result is None
 
 
-class TestMessageExistsInFolder:
-    """Tests for message_exists_in_folder function."""
-
-    def test_no_message_id(self):
-        """Test returns False when message_id is None."""
-        mock_conn = Mock()
-        result = imap_common.message_exists_in_folder(mock_conn, None)
-        assert result is False
-
-    def test_search_fails(self):
-        """Test returns False when search fails."""
-        mock_conn = Mock()
-        mock_conn.search.return_value = ("NO", [])
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is False
-
-    def test_no_matches(self):
-        """Test returns False when no matches found."""
-        mock_conn = Mock()
-        mock_conn.search.return_value = ("OK", [b""])
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is False
-
-    def test_match_found(self):
-        """Test returns True when message with same ID found."""
-        mock_conn = Mock()
-        mock_conn.search.return_value = ("OK", [b"1"])
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is True
-
-    def test_search_exception(self):
-        """Test returns False when search raises an exception."""
-        mock_conn = Mock()
-        mock_conn.search.side_effect = Exception("Connection error")
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is False
-
-
 class TestExtractMessageId:
     """Tests for extract_message_id function."""
 
@@ -594,20 +604,21 @@ class TestAppendEmail:
 
         assert result is False
 
-    def test_append_exception_returns_false(self):
-        """Test append returns False when exception occurs."""
+    def test_append_exception_propagates(self):
+        """Test append propagates exceptions so callers can handle/log."""
         mock_conn = Mock()
         mock_conn.append.side_effect = Exception("Connection error")
 
-        result = imap_common.append_email(
-            mock_conn,
-            "TestFolder",
-            b"email content",
-            "01-Jan-2024 12:00:00 +0000",
-            ensure_folder=False,
-        )
+        import pytest
 
-        assert result is False
+        with pytest.raises(Exception, match="Connection error"):
+            imap_common.append_email(
+                mock_conn,
+                "TestFolder",
+                b"email content",
+                "01-Jan-2024 12:00:00 +0000",
+                ensure_folder=False,
+            )
 
     def test_append_with_multiple_flags(self):
         """Test append with multiple flags."""
@@ -652,14 +663,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()
@@ -706,8 +725,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"]),
@@ -717,6 +736,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()
@@ -735,3 +765,77 @@ class TestGetMessageIdsInFolder:
 
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert set(result.values()) == {"<valid@example.com>"}
+
+
+class TestLoadFolderMsgIds:
+    """Tests for load_folder_msg_ids() using the mock IMAP server."""
+
+    def test_returns_none_when_disabled(self):
+        """Returns None if dict or lock is None (tracking disabled)."""
+        mock_conn = Mock()
+        import threading
+
+        lock = threading.Lock()
+        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, lock) is None
+        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", {}, None) is None
+        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, None) is None
+
+    def test_fetches_message_ids_from_server(self, single_mock_server):
+        """Fetches Message-IDs from server and caches them."""
+        import imaplib
+        import threading
+
+        server_data = {
+            "TestFolder": [
+                b"Subject: A\r\nMessage-ID: <a@test.com>\r\n\r\nBody A",
+                b"Subject: B\r\nMessage-ID: <b@test.com>\r\n\r\nBody B",
+            ]
+        }
+        _server, port = single_mock_server(server_data)
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        lock = threading.Lock()
+        by_folder: dict[str, set[str]] = {}
+
+        result = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
+        assert result == {"<a@test.com>", "<b@test.com>"}
+        assert "TestFolder" in by_folder
+
+        # Second call returns the same cached set object
+        result2 = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
+        assert result2 is result
+        conn.logout()
+
+    def test_empty_folder(self, single_mock_server):
+        """Returns empty set for a folder with no messages."""
+        import imaplib
+        import threading
+
+        _server, port = single_mock_server({"INBOX": []})
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        lock = threading.Lock()
+        by_folder: dict[str, set[str]] = {}
+
+        result = imap_common.load_folder_msg_ids(conn, "INBOX", by_folder, lock)
+        assert result == set()
+        conn.logout()
+
+    def test_creates_folder_if_missing(self, single_mock_server):
+        """Creates folder on first access if it doesn't exist."""
+        import imaplib
+        import threading
+
+        server, port = single_mock_server({"INBOX": []})
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        lock = threading.Lock()
+        by_folder: dict[str, set[str]] = {}
+
+        result = imap_common.load_folder_msg_ids(conn, "NewFolder", by_folder, lock)
+        assert result == set()
+        assert "NewFolder" in server.folders
+        conn.logout()

test/utils/test_imap_compress.py

@@ -0,0 +1,394 @@
+"""Tests for IMAP COMPRESS=DEFLATE support."""
+
+import zlib
+from unittest.mock import MagicMock, patch
+
+from utils import imap_compress
+
+
+class TestCompressedSocket:
+    """Tests for the _CompressedSocket wrapper."""
+
+    def _make_compressed(self, plaintext: bytes) -> bytes:
+        """Compress data using raw DEFLATE with Z_SYNC_FLUSH (server-side)."""
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        return c.compress(plaintext) + c.flush(zlib.Z_SYNC_FLUSH)
+
+    def test_sendall_compresses_data(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+
+        ws.sendall(b"AAAA0 SELECT INBOX\r\n")
+
+        sock.sendall.assert_called_once()
+        compressed = sock.sendall.call_args[0][0]
+        # Verify round-trip: decompress the sent data
+        d = zlib.decompressobj(-zlib.MAX_WBITS)
+        assert d.decompress(compressed) == b"AAAA0 SELECT INBOX\r\n"
+
+    def test_send_compresses_and_returns_length(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+
+        result = ws.send(b"hello")
+        assert result == 5
+        sock.sendall.assert_called_once()
+
+    def test_recv_decompresses_data(self):
+        plaintext = b"* OK IMAP ready\r\n"
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        sock.recv.return_value = compressed
+
+        ws = imap_compress._CompressedSocket(sock)
+        result = ws.recv(4096)
+
+        assert result == plaintext
+
+    def test_recv_buffers_excess(self):
+        plaintext = b"A" * 100
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        sock.recv.return_value = compressed
+
+        ws = imap_compress._CompressedSocket(sock)
+        # Request only 30 bytes — rest should be buffered
+        chunk1 = ws.recv(30)
+        assert chunk1 == b"A" * 30
+        assert len(ws._recv_buf) == 70
+
+        # Next recv should return from buffer without hitting socket
+        chunk2 = ws.recv(70)
+        assert chunk2 == b"A" * 70
+        sock.recv.assert_called_once()  # Only one real socket read
+
+    def test_recv_empty_returns_empty(self):
+        sock = MagicMock()
+        sock.recv.return_value = b""
+
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.recv(4096) == b""
+
+    def test_initial_compressed_data_decompressed(self):
+        plaintext = b"* 5 EXISTS\r\n"
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock, initial_compressed=compressed)
+
+        # Should return the pre-decompressed data without hitting socket
+        result = ws.recv(4096)
+        assert result == plaintext
+        sock.recv.assert_not_called()
+
+    def test_shutdown_proxied(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+        ws.shutdown(2)
+        sock.shutdown.assert_called_once_with(2)
+
+    def test_close_proxied(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+        ws.close()
+        sock.close.assert_called_once()
+
+    def test_makefile_returns_buffered_reader(self):
+        import io
+
+        plaintext = b"* OK ready\r\n"
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        sock.recv.return_value = compressed
+
+        ws = imap_compress._CompressedSocket(sock)
+        f = ws.makefile("rb")
+        assert isinstance(f, io.BufferedReader)
+        assert f.readline() == plaintext
+
+    def test_getattr_proxied(self):
+        sock = MagicMock()
+        sock.gettimeout.return_value = 30
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.gettimeout() == 30
+
+    def test_multiple_sendall_share_compressor_state(self):
+        """Verify the compressor maintains state across calls (streaming)."""
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+
+        ws.sendall(b"first ")
+        ws.sendall(b"second")
+
+        assert sock.sendall.call_count == 2
+        # Both calls should produce decompressible data
+        d = zlib.decompressobj(-zlib.MAX_WBITS)
+        result = b""
+        for call in sock.sendall.call_args_list:
+            result += d.decompress(call[0][0])
+        assert result == b"first second"
+
+    def test_recv_loops_on_partial_deflate_block(self):
+        """Decompress can return b'' for partial blocks; recv must not return that as EOF."""
+        sock = MagicMock()
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+
+        # Build a compressed stream, then split it mid-block so the first
+        # chunk decompresses to b"" and the second completes the block.
+        full_compressed = c.compress(b"hello\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+        # Split at byte 1 — the first fragment is too short to produce output
+        part1 = full_compressed[:1]
+        part2 = full_compressed[1:]
+        sock.recv.side_effect = [part1, part2]
+
+        ws = imap_compress._CompressedSocket(sock)
+        result = ws.recv(4096)
+
+        assert result == b"hello\r\n"
+        assert sock.recv.call_count == 2
+
+    def test_recv_returns_empty_on_real_eof(self):
+        """True socket EOF (b'') should still return b''."""
+        sock = MagicMock()
+        sock.recv.return_value = b""
+
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.recv(4096) == b""
+
+    def test_recv_multiple_chunks(self):
+        """Simulate receiving data across multiple socket reads."""
+        chunk1 = b"line one\r\n"
+        chunk2 = b"line two\r\n"
+
+        sock = MagicMock()
+        # Server compressor maintains state across writes
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed1 = c.compress(chunk1) + c.flush(zlib.Z_SYNC_FLUSH)
+        compressed2 = c.compress(chunk2) + c.flush(zlib.Z_SYNC_FLUSH)
+        sock.recv.side_effect = [compressed1, compressed2]
+
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.recv(4096) == chunk1
+        assert ws.recv(4096) == chunk2
+
+
+class TestEnableCompression:
+    """Tests for the enable_compression() function."""
+
+    def test_returns_false_when_capability_missing(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "IDLE")
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_capability_missing_bytes(self):
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"IDLE")
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_no_capabilities_attr(self):
+        conn = MagicMock(spec=[])
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_command_fails(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("NO", [b"compression not available"])
+
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_command_raises(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.side_effect = Exception("network error")
+
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_enables_compression_with_bytes_capabilities(self):
+        """Real imaplib stores capabilities as bytes; verify they are handled."""
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"DEFLATE active"])
+        conn._readbuf = []
+        conn.sock = MagicMock()
+
+        result = imap_compress.enable_compression(conn)
+
+        assert result is True
+        assert isinstance(conn.sock, imap_compress._CompressedSocket)
+
+    def test_enables_compression_on_success(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"DEFLATE active"])
+        conn._readbuf = []
+        conn.sock = MagicMock()
+
+        result = imap_compress.enable_compression(conn)
+
+        assert result is True
+        assert isinstance(conn.sock, imap_compress._CompressedSocket)
+
+    def test_drains_readbuf_on_success(self):
+        # Simulate leftover compressed data in the read buffer
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 1 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = [compressed]
+        conn.sock = MagicMock()
+
+        imap_compress.enable_compression(conn)
+
+        # readbuf should have been cleared
+        assert len(conn._readbuf) == 0
+        # The compressed leftover should be pre-decompressed in the wrapper
+        assert conn.sock._recv_buf == b"* 1 EXISTS\r\n"
+
+    def test_drains_bytearray_readbuf(self):
+        """Real imaplib uses bytearray for _readbuf; verify drain handles it."""
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 2 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = bytearray(compressed)
+        conn.sock = MagicMock()
+
+        imap_compress.enable_compression(conn)
+
+        assert len(conn._readbuf) == 0
+        assert conn.sock._recv_buf == b"* 2 EXISTS\r\n"
+
+    def test_drains_bytes_readbuf(self):
+        """Some imaplib versions use bytes for _readbuf; verify drain handles it."""
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 3 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = bytes(compressed)
+        conn.sock = MagicMock()
+
+        imap_compress.enable_compression(conn)
+
+        assert conn._readbuf == b""
+        assert conn.sock._recv_buf == b"* 3 EXISTS\r\n"
+
+    def test_drains_file_buffer(self):
+        """BufferedReader may have read-ahead bytes beyond the OK response."""
+        import io
+
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 4 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = b""
+        conn.sock = MagicMock()
+
+        # Simulate a BufferedReader with read-ahead compressed bytes
+        old_file = io.BufferedReader(io.BytesIO(compressed))
+        old_file.peek(65536)  # fill the buffer
+        conn._file = old_file
+
+        imap_compress.enable_compression(conn)
+
+        assert conn.sock._recv_buf == b"* 4 EXISTS\r\n"
+
+    def test_calls_log_fn(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = []
+        conn.sock = MagicMock()
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "compression enabled" in messages[0]
+
+    def test_logs_when_not_supported(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1",)  # No COMPRESS
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "not supported" in messages[0]
+
+    def test_logs_when_rejected(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("NO", [b"not available"])
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "rejected" in messages[0]
+
+    def test_logs_when_negotiation_fails(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.side_effect = Exception("network error")
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "failed" in messages[0]
+
+    def test_recreates_file(self):
+        import io
+
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = []
+        real_sock = MagicMock()
+        conn.sock = real_sock
+
+        imap_compress.enable_compression(conn)
+
+        # _file should be a BufferedReader over the compressed socket
+        assert isinstance(conn._file, io.BufferedReader)
+
+
+class TestIntegrationWithGetImapConnection:
+    """Test that enable_compression is called during connection setup."""
+
+    @patch("utils.imap_common.imaplib")
+    @patch("utils.imap_compress.enable_compression")
+    def test_compression_attempted_after_login(self, mock_compress, mock_imaplib):
+        from utils import imap_common
+
+        mock_conn = MagicMock()
+        mock_imaplib.IMAP4_SSL.return_value = mock_conn
+
+        imap_common.get_imap_connection("imap.example.com", "user", password="pass")
+
+        mock_conn.login.assert_called_once()
+        mock_compress.assert_called_once_with(mock_conn, log_fn=imap_common.safe_print)
+
+    @patch("utils.imap_common.imaplib")
+    @patch("utils.imap_compress.enable_compression")
+    def test_compression_attempted_after_oauth2(self, mock_compress, mock_imaplib):
+        from utils import imap_common
+
+        mock_conn = MagicMock()
+        mock_imaplib.IMAP4_SSL.return_value = mock_conn
+
+        imap_common.get_imap_connection("imap.example.com", "user", oauth2_token="token123")
+
+        mock_conn.authenticate.assert_called_once()
+        mock_compress.assert_called_once_with(mock_conn, log_fn=imap_common.safe_print)

tools/mock_imap_server.py

@@ -1,4 +1,3 @@
-import re
 import socketserver
 import threading
 
@@ -15,7 +14,15 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
     def handle(self):
         self.wfile.write(b"* OK [CAPABILITY IMAP4rev1] Mock IMAP Server Ready\r\n")
         self.selected_folder = None
-        self.current_folders = self.server.folders
+
+        # Check if server has folders, if not, initialize
+        if hasattr(self.server, "folders"):
+            self.current_folders = self.server.folders
+        else:
+            # Just in case it's used standalone without the threaded wrapper
+            self.current_folders = {"INBOX": []}
+
+        self.retry_state = {}  # key -> count
 
         while True:
             try:
@@ -31,9 +38,73 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 cmd = parts[1].upper()
                 args = parts[2] if len(parts) > 2 else ""
 
+                # Check for RETRY_N injection in args (for testing imap_retry)
+                # Pattern: RETRY_N where N is integer
+                import re
+
+                retry_match = re.search(r"RETRY_(\d+)", args)
+                if retry_match:
+                    count = int(retry_match.group(1))
+                    # Create a unique key for this command invocation context
+                    # Just using command + args is roughly sufficient
+                    key = f"{cmd}_{args}"
+                    current_fails = self.retry_state.get(key, 0)
+
+                    if current_fails < count:
+                        # Logic Fix: Only increment if we are going to fail
+                        self.retry_state[key] = current_fails + 1
+                        self.send_response(tag, "NO [UNAVAILABLE] Server Busy (Simulated)")
+                        continue
+                    # Else fall through to normal command processing
+
                 if cmd == "LOGIN":
                     self.send_response(tag, "OK LOGIN completed")
 
+                elif cmd == "APPEND":
+                    # Check for literal size
+                    import re
+
+                    size_match = re.search(r"\{(\d+)\}$", args)
+                    data = b""
+                    if size_match:
+                        size = int(size_match.group(1))
+                        # Send continuation
+                        self.wfile.write(b"+\r\n")
+                        self.wfile.flush()
+                        # Read data
+                        data = self.rfile.read(size)
+
+                    # Extract mailbox name to store the message
+                    # Simplistic parsing: check for quoted or unquoted first argument
+                    mailbox = "INBOX"  # Default fallback
+                    clean_args = args.lstrip()
+                    if clean_args.startswith('"'):
+                        end_quote = clean_args.find('"', 1)
+                        if end_quote != -1:
+                            mailbox = clean_args[1:end_quote]
+                    else:
+                        mailbox = clean_args.split(" ")[0]
+
+                    if mailbox not in self.current_folders:
+                        self.current_folders[mailbox] = []
+
+                    # Parse flags if present: look for (...)
+                    msg_flags = set()
+                    flag_match = re.search(r"\(([^)]+)\)", args)
+                    if flag_match:
+                        # e.g. (\Seen \Deleted)
+                        flag_content = flag_match.group(1)
+                        msg_flags = set(flag_content.split())
+
+                    # Store message
+                    new_uid = len(self.current_folders[mailbox]) + 1
+                    msg_obj = {"uid": new_uid, "flags": msg_flags, "content": data}
+                    self.current_folders[mailbox].append(msg_obj)
+
+                    # Always succeed for test if we passed retry logic
+                    self.wfile.write(tag.encode() + b" OK [APPENDUID 1 100] APPEND completed\r\n")
+                    self.wfile.flush()
+
                 elif cmd == "LOGOUT":
                     self.send_response(tag, "OK LOGOUT completed")
                     break
@@ -42,6 +113,16 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                     self.wfile.write(b"* CAPABILITY IMAP4rev1 AUTH=PLAIN\r\n")
                     self.send_response(tag, "OK CAPABILITY completed")
 
+                elif cmd == "AUTHENTICATE":
+                    # Minimal XOAUTH2 support for tests.
+                    self.wfile.write(b"+ \r\n")
+                    _ = self.rfile.readline()
+                    # Also consume potential retry line? No.
+                    self.send_response(tag, "OK AUTHENTICATE completed")
+
+                elif cmd == "NOOP":
+                    self.send_response(tag, "OK NOOP completed")
+
                 elif cmd == "LIST":
                     for folder in self.current_folders:
                         self.wfile.write(f'* LIST (\\HasNoChildren) "/" "{folder}"\r\n'.encode())
@@ -49,23 +130,40 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
 
                 elif cmd == "SELECT":
                     folder = args.strip().strip('"')
-                    if folder in self.current_folders:
+
+                    # For retry testing: if folder starts with NO, fail immediately
+                    if folder.startswith("NO"):
+                        # e.g. "NO [UNAVAILABLE]"
+                        self.send_response(tag, folder)
+                        continue
+
+                    # Allow RETRY_N folders to succeed if they passed the retry intercept logic
+                    is_retry_folder = "RETRY_" in folder
+
+                    if folder in self.current_folders or is_retry_folder:
                         self.selected_folder = folder
-                        count = len(self.current_folders[folder])
+                        count = len(self.current_folders.get(folder, []))  # Default empty for retry folders
                         self.wfile.write(f"* {count} EXISTS\r\n".encode())
                         self.wfile.write(f"* {count} RECENT\r\n".encode())
                         self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
                         self.wfile.write(b"* OK [UIDVALIDITY 1] UIDs valid\r\n")
                         self.send_response(tag, "OK [READ-WRITE] SELECT completed")
+                    elif folder == "EMPTY":
+                        # Special case for "empty_data" test in retry
+                        self.wfile.write(b"* 0 EXISTS\r\n")
+                        self.send_response(tag, "OK SELECT completed")
                     else:
                         self.send_response(tag, "NO [NONEXISTENT] Folder not found")
 
                 elif cmd == "EXAMINE":
                     # EXAMINE is like SELECT but read-only
                     folder = args.strip().strip('"')
-                    if folder in self.current_folders:
+
+                    is_retry_folder = "RETRY_" in folder
+
+                    if folder in self.current_folders or is_retry_folder:
                         self.selected_folder = folder
-                        count = len(self.current_folders[folder])
+                        count = len(self.current_folders.get(folder, []))
                         self.wfile.write(f"* {count} EXISTS\r\n".encode())
                         self.wfile.write(f"* {count} RECENT\r\n".encode())
                         self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
@@ -80,6 +178,42 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                         self.current_folders[folder] = []
                     self.send_response(tag, "OK CREATE completed")
 
+                elif cmd == "SEARCH":
+                    if not self.selected_folder:
+                        self.send_response(tag, RESPONSE_SELECT_FIRST)
+                        continue
+
+                    msgs = self.current_folders[self.selected_folder]
+                    sub_args = args
+
+                    header_msg_id = None
+                    try:
+                        m = re.search(r'HEADER\s+Message-ID\s+"([^"]+)"', sub_args, re.IGNORECASE)
+                        if m:
+                            header_msg_id = m.group(1)
+                    except Exception:
+                        header_msg_id = None
+
+                    seq_nums = []
+                    for idx, m in enumerate(msgs, start=1):
+                        if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]:
+                            continue
+                        if header_msg_id:
+                            msg_text = m["content"].decode("utf-8", errors="ignore")
+                            if header_msg_id not in msg_text:
+                                continue
+                        seq_nums.append(str(idx))
+
+                    if "ALL" in sub_args.upper() and not header_msg_id:
+                        seq_nums = [str(idx) for idx in range(1, len(msgs) + 1)]
+
+                    seq_str = " ".join(seq_nums)
+                    if seq_str:
+                        self.wfile.write(f"* SEARCH {seq_str}\r\n".encode())
+                    else:
+                        self.wfile.write(b"* SEARCH\r\n")
+                    self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
+
                 elif cmd == "EXPUNGE":
                     if self.selected_folder:
                         msgs = self.current_folders[self.selected_folder]
@@ -125,7 +259,10 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                             valid_uids.append(str(m["uid"]))
 
                         uids_str = " ".join(valid_uids)
-                        self.wfile.write(f"* SEARCH {uids_str}\r\n".encode())
+                        if uids_str:
+                            self.wfile.write(f"* SEARCH {uids_str}\r\n".encode())
+                        else:
+                            self.wfile.write(b"* SEARCH\r\n")
                         self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
 
                     elif sub_cmd == "STORE":
@@ -292,50 +429,6 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                         print(f"MOCK APPEND ERROR: {e}")
                         self.send_response(tag, "BAD APPEND")
 
-                elif cmd == "SEARCH":
-                    # Parse SEARCH ALL, SEARCH HEADER Message-ID "...", etc.
-                    # args might be: ALL, HEADER Message-ID "<123>", CHARSETS UTF-8 ...
-
-                    found_indices = []
-
-                    if not self.selected_folder:
-                        self.wfile.write(b"* SEARCH\r\n")
-                        self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
-                        continue
-
-                    msgs = self.current_folders[self.selected_folder]
-
-                    if "ALL" in args.upper():
-                        # Return all message sequence numbers
-                        found_indices = [str(idx + 1) for idx in range(len(msgs))]
-                    elif "HEADER Message-ID" in args:
-                        # Extract value
-                        # Expected: ... HEADER Message-ID "value" ...
-                        try:
-                            # Split by 'MESSAGE-ID' (case insensitive?)
-                            # part after Message-ID
-                            post_mi = args.split("Message-ID", 1)[1].strip()
-                            # Should start with quote or value
-                            if post_mi.startswith('"'):
-                                search_val = post_mi.split('"', 2)[1]
-                            else:
-                                search_val = post_mi.split(" ", 1)[0]
-
-                            search_val = search_val.replace("<", "").replace(">", "")
-
-                            for idx, m in enumerate(msgs):
-                                content_str = m["content"].decode("utf-8", errors="ignore")
-                                if search_val in content_str:
-                                    # Simple substring check is risky but okay for mock
-                                    # Better: regex for Message-ID: <...search_val...>
-                                    found_indices.append(str(idx + 1))
-                        except Exception:
-                            pass
-
-                    indices_str = " ".join(found_indices)
-                    self.wfile.write(f"* SEARCH {indices_str}\r\n".encode())
-                    self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
-
                 elif cmd == "STORE":
                     # STORE <msg_set> +FLAGS (\Seen)  (non-UID; uses message sequence numbers)
                     if not self.selected_folder:
@@ -411,7 +504,11 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 else:
                     self.send_response(tag, "BAD Command not recognized")
 
-            except Exception:
+            except Exception as e:
+                print(f"MockServer EXCEPTION: {e}")
+                import traceback
+
+                traceback.print_exc()
                 break
 
     def send_response(self, tag, message):
@@ -437,9 +534,19 @@ class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
             self.folders = {"INBOX": []}
 
 
-def start_server_thread(port=10143, initial_folders=None):
+def start_server_thread(port=0, initial_folders=None):
+    # Use port 0 to let OS select free port
+    if isinstance(port, dict):
+        # Handle legacy call where port was inadvertently passed as dict
+        initial_folders = port
+        port = 0
+
     server = MockIMAPServer(("localhost", port), MockIMAPHandler, initial_folders)
+
+    # Get the actual port if 0 was used
+    actual_port = server.server_address[1]
+
     t = threading.Thread(target=server.serve_forever)
     t.daemon = True
     t.start()
-    return t, server
+    return server, actual_port

tools/mock_oauth_server.py

@@ -0,0 +1,59 @@
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from urllib.parse import urlparse
+
+MOCK_TENANT_ID = "00000000-0000-0000-0000-000000000000"
+
+
+def _write_json(handler, status, payload):
+    body = json.dumps(payload).encode("utf-8")
+    handler.send_response(status)
+    handler.send_header("Content-Type", "application/json")
+    handler.send_header("Content-Length", str(len(body)))
+    handler.end_headers()
+    handler.wfile.write(body)
+
+
+class MockOAuthHandler(BaseHTTPRequestHandler):
+    """Minimal OAuth2 mock server for tests."""
+
+    def do_GET(self):
+        parsed = urlparse(self.path)
+        if parsed.path.endswith("/.well-known/openid-configuration"):
+            issuer = f"https://login.microsoftonline.com/{MOCK_TENANT_ID}/v2.0"
+            _write_json(self, 200, {"issuer": issuer})
+            return
+
+        _write_json(self, 404, {"error": "not_found"})
+
+    def do_POST(self):
+        parsed = urlparse(self.path)
+        length = int(self.headers.get("Content-Length", "0"))
+        _ = self.rfile.read(length) if length else b""
+
+        if parsed.path == "/google/token":
+            _write_json(self, 200, {"access_token": "mock-google-token"})
+            return
+
+        if parsed.path == "/microsoft/token":
+            _write_json(self, 200, {"access_token": "mock-microsoft-token"})
+            return
+
+        _write_json(self, 404, {"error": "not_found"})
+
+    def log_message(self, _format, *_args):
+        # Silence default HTTP server logging during tests.
+        return
+
+
+class MockOAuthServer(HTTPServer):
+    allow_reuse_address = True
+
+
+def start_server_thread(port=0):
+    server = MockOAuthServer(("localhost", port), MockOAuthHandler)
+    thread = threading.Thread(target=server.serve_forever)
+    thread.daemon = True
+    thread.start()
+    return thread, server