Initial plan

* Initial plan

* Add comprehensive unit tests for ensure_folder_exists and append_email

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

.github/copilot-instructions.md

@@ -1,8 +0,0 @@
-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 utils.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 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

@@ -1,59 +0,0 @@
-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,4 +7,3 @@ 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. **`imap_migrate.py`** (The Solution)
+1. **`migrate_imap_emails.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,20 +25,19 @@ 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. **`imap_compare.py`** (The Validator)
+2. **`compare_imap_folders.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. **`imap_count.py`** (The Investigator)
+3. **`count_imap_emails.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. **`imap_backup.py`** (The Backup)
+4. **`backup_imap_emails.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.
@@ -46,7 +45,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. **`imap_restore.py`** (The Restore)
+5. **`restore_imap_emails.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.
@@ -67,36 +66,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### 2. Installation
 
-#### 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.
+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.
@@ -158,7 +128,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```bash
-   python3 imap_migrate.py
+   python3 migrate_imap_emails.py
    ```
 
 #### Windows (PowerShell)
@@ -175,7 +145,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```powershell
-   python imap_migrate.py
+   python migrate_imap_emails.py
    ```
 
 ### Method 2: Command Line Arguments (Overrides)
@@ -183,7 +153,7 @@ All scripts support command-line arguments which take precedence over environmen
 
 **Migration:**
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -196,7 +166,7 @@ python3 imap_migrate.py \
 
 **Comparison:**
 ```bash
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -205,14 +175,14 @@ python3 imap_compare.py \
   --dest-pass "your-app-password"
 
 # Compare IMAP source to a local backup folder
-python3 imap_compare.py \
+python3 compare_imap_folders.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 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -221,23 +191,23 @@ python3 imap_compare.py \
 
 **Counting:**
 ```bash
-python3 imap_count.py \
+python3 count_imap_emails.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --pass "secret"
 
 # Count a local backup folder
-python3 imap_count.py \
+python3 count_imap_emails.py \
    --path "./my_backup"
 
 # Or via environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 imap_count.py
+python3 count_imap_emails.py
 ```
 
 **Counting (OAuth2):**
 ```bash
-python3 imap_count.py \
+python3 count_imap_emails.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --oauth2-client-id "id" \
@@ -248,12 +218,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 imap_count.py
+python3 count_imap_emails.py
 ```
 
 **Backup:**
 ```bash
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -265,7 +235,7 @@ python3 imap_backup.py \
 ### 1. Full Migration
 Migrate all folders from Source to Destination.
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -278,7 +248,7 @@ python3 imap_migrate.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 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --gmail-mode \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
@@ -288,29 +258,17 @@ python3 imap_migrate.py \
   --dest-pass "dest-app-password"
 ```
 
-### 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.
+Notes:
+- `--preserve-flags` is enabled automatically in `--gmail-mode`.
+- `--dest-delete` is not supported in `--gmail-mode`.
 
-```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)
+### 1b. 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 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --preserve-flags \
   --src-host "imap.example.com" \
   --src-user "source@example.com" \
@@ -323,8 +281,8 @@ python3 imap_migrate.py \
 ### 2. Single Folder Migration
 Migrate ONLY a specific folder (e.g., trying to fix just "Important" or "Sent").
 ```bash
-# Syntax: python3 imap_migrate.py "[Folder Name]"
-python3 imap_migrate.py \
+# Syntax: python3 migrate_imap_emails.py "[Folder Name]"
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -338,7 +296,7 @@ python3 imap_migrate.py \
 Migrate and **delete** from source immediately after verifying the copy.
 ```bash
 # Using flag
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -348,7 +306,7 @@ python3 imap_migrate.py \
   --src-delete
 
 # Or specific folder with delete
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -365,7 +323,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 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -375,7 +333,7 @@ python3 imap_migrate.py \
   --dest-delete
 
 # Backup: Delete local .eml files not found on server
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -383,7 +341,7 @@ python3 imap_backup.py \
   --dest-delete
 
 # Restore: Delete server emails not found in local backup
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -396,7 +354,7 @@ python3 imap_restore.py \
 ### 5. Verify Migration
 Compare counts between source and destination.
 ```bash
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -405,14 +363,14 @@ python3 imap_compare.py \
   --dest-pass "dest-app-password"
 
 # IMAP source -> local backup destination
-python3 imap_compare.py \
+python3 compare_imap_folders.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 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -430,21 +388,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 imap_backup.py \
+python3 backup_imap_emails.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 imap_backup.py \
+python3 backup_imap_emails.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 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -453,18 +411,18 @@ python3 imap_backup.py \
 ```
 
 ### 6a. Compare IMAP vs Local Backup
-Use `imap_compare.py` to validate an IMAP account against a local backup created by `imap_backup.py`.
+Use `compare_imap_folders.py` to validate an IMAP account against a local backup created by `backup_imap_emails.py`.
 
 ```bash
 # Option 1: IMAP source -> local destination
-python3 imap_compare.py \
+python3 compare_imap_folders.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 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -479,15 +437,15 @@ export DEST_LOCAL_PATH="./my_backup"
 ```
 
 ### 6b. Count a Local Backup
-Use `imap_count.py` to get per-folder counts from a local backup created by `imap_backup.py`.
+Use `count_imap_emails.py` to get per-folder counts from a local backup created by `backup_imap_emails.py`.
 
 ```bash
 # Option 1: explicit path
-python3 imap_count.py --path "./my_backup"
+python3 count_imap_emails.py --path "./my_backup"
 
 # Option 2: environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 imap_count.py
+python3 count_imap_emails.py
 ```
 
 ### 7. Gmail Backup with Labels Preservation
@@ -495,7 +453,7 @@ When backing up a Gmail account, use `--gmail-mode` for the recommended workflow
 
 ```bash
 # Recommended: Use --gmail-mode for simplest workflow
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -505,7 +463,7 @@ python3 imap_backup.py \
 
 This is equivalent to the more verbose:
 ```bash
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -517,7 +475,7 @@ python3 imap_backup.py \
 **For large accounts (100K+ emails)**, you can build the manifest first to test:
 ```bash
 # Step 1: Build manifest only (fast, no download)
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -525,7 +483,7 @@ python3 imap_backup.py \
   --manifest-only
 
 # Step 2: Download emails (can run later, manifest already exists)
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -571,7 +529,7 @@ python3 imap_backup.py \
 For non-Gmail servers, you can preserve read/starred status with `--preserve-flags`:
 
 ```bash
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.example.com" \
   --src-user "you@example.com" \
   --src-pass "your-password" \
@@ -590,14 +548,14 @@ Use `--full-restore` to force the legacy behavior (process all emails and re-syn
 
 ```bash
 # Restore all folders from backup
-python3 imap_restore.py \
+python3 restore_imap_emails.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 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -605,7 +563,7 @@ python3 imap_restore.py \
   --full-restore
 
 # Restore with flags (read/starred status)
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.example.com" \
   --dest-user "you@example.com" \
@@ -613,7 +571,7 @@ python3 imap_restore.py \
   --apply-flags
 
 # Restore a specific folder
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -625,7 +583,7 @@ python3 imap_restore.py \
 Restore a Gmail backup with full label structure using `--gmail-mode`:
 
 ```bash
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -642,7 +600,7 @@ python3 imap_restore.py \
 
 **Alternatively**, restore folders individually with labels and flags applied:
 ```bash
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -664,7 +622,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 `imap_count.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_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`).
 
 ### Microsoft (Outlook / Office 365)
 
@@ -703,7 +661,7 @@ Requires the `msal` package (`pip install msal`). Uses the **device code flow**
 pip install msal
 
 # Migration with Microsoft OAuth2 on source
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "outlook.office365.com" \
   --src-user "user@contoso.com" \
   --src-oauth2-client-id "your-azure-app-client-id" \
@@ -723,7 +681,7 @@ Requires the `google-auth-oauthlib` package (`pip install google-auth-oauthlib`)
 pip install google-auth-oauthlib
 
 # Backup with Google OAuth2
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-oauth2-client-id "your-google-client-id" \
@@ -736,7 +694,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 `imap_migrate.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 `migrate_imap_emails.py` uses multiple threads, you may hit this limit.
   **Solution**: Reduce `MAX_WORKERS` to `4` or `2` using the environment variable.
 
 - **Authentication Errors**:
@@ -789,7 +747,7 @@ PYTHONPATH=src pytest test/ -v
 make coverage
 
 # Run a specific test file
-PYTHONPATH=src pytest test/test_imap_migrate.py -v
+PYTHONPATH=src pytest test/test_migrate_imap_emails.py -v
 
 # Run a specific test
 PYTHONPATH=src pytest test/test_imap_common.py::TestNormalizeFolderName -v
@@ -821,11 +779,11 @@ make ci
 
 | Test File | Description |
 |-----------|-------------|
-| `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_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_common.py` | Shared utility function tests |
 
 ### Continuous Integration

pyproject.toml

@@ -1,44 +1,6 @@
 # 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

@@ -1,201 +0,0 @@
-"""
-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

@@ -1,72 +0,0 @@
-"""
-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

@@ -1,154 +0,0 @@
-"""
-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,24 +1,894 @@
-#!/usr/bin/env python3
 """
-DEPRECATED: This script has been renamed.
-Please use imap_backup.py instead.
+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.
 """
 
+import argparse
+import concurrent.futures
+import json
+import os
 import sys
+import threading
 
-from imap_backup import main
+import imap_common
+import imap_oauth2
 
-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)
+# 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)
     try:
-        main()
+        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()
+
     except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
+        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.")
         sys.exit(0)
     except Exception as e:
         print(f"Fatal Error: {e}")
-        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

src/compare_imap_folders.py

@@ -1,24 +1,406 @@
-#!/usr/bin/env python3
 """
-DEPRECATED: This script has been renamed.
-Please use imap_compare.py instead.
+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"
 """
 
+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__":
-    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)
+    main()

src/core/imap_session.py

@@ -1,123 +0,0 @@
-"""
-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,24 +1,296 @@
-#!/usr/bin/env python3
-"""
-DEPRECATED: This script has been renamed.
-Please use imap_count.py instead.
+"""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
 """
 
+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__":
-    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)
+    main()

src/imap_backup.py

@@ -1,899 +0,0 @@
-"""
-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

@@ -0,0 +1,438 @@
+"""
+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

@@ -1,337 +0,0 @@
-"""
-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

@@ -1,247 +0,0 @@
-"""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

@@ -0,0 +1,258 @@
+"""
+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

@@ -1,976 +0,0 @@
-"""
-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

@@ -1,29 +0,0 @@
-"""
-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

@@ -1,130 +0,0 @@
-"""
-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/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) -> bool:
+def save_dest_index_cache(cache_path: str, cache_data: dict) -> None:
     try:
         tmp_path = f"{cache_path}.tmp"
         with open(tmp_path, "w", encoding="utf-8") as f:
@@ -252,62 +252,3 @@ 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,
-        )

src/utils/imap_common.py

@@ -1,785 +0,0 @@
-"""
-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

@@ -1,180 +0,0 @@
-"""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

@@ -1,80 +0,0 @@
-"""
-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

test/auth/test_imap_oauth2.py

@@ -1,416 +0,0 @@
-"""
-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

@@ -1,169 +0,0 @@
-"""
-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

@@ -1,290 +0,0 @@
-"""
-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,7 +7,6 @@ import os
 import socket
 import sys
 import time
-from contextlib import contextmanager
 
 import pytest
 
@@ -16,7 +15,6 @@ 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():
@@ -33,19 +31,11 @@ def create_server_pair(src_data=None, dest_data=None):
     while p2 == p1:
         p2 = get_free_port()
 
-    src_s, p1_actual = start_server_thread(p1, src_data)
-    dest_s, p2_actual = start_server_thread(p2, dest_data)
+    src_t, src_s = start_server_thread(p1, src_data)
+    dest_t, dest_s = start_server_thread(p2, dest_data)
     time.sleep(0.3)
 
-    # 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)
+    return (src_t, src_s, p1), (dest_t, dest_s, p2)
 
 
 def shutdown_server_pair(src_tuple, dest_tuple):
@@ -53,13 +43,9 @@ 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()
-    dest_s.server_close()
-    if src_t:
-        src_t.join(timeout=2)
-    if dest_t:
-        dest_t.join(timeout=2)
+    src_t.join(timeout=2)
+    dest_t.join(timeout=2)
 
 
 @pytest.fixture
@@ -94,47 +80,22 @@ def single_mock_server():
 
     def _create(initial_data=None):
         port = get_free_port()
-        server, actual_port = start_server_thread(port, initial_data)
+        thread, server = start_server_thread(port, initial_data)
         time.sleep(0.3)
-        servers.append(server)
-        return server, actual_port
+        servers.append((thread, server))
+        return server, port
 
     yield _create
 
-    for server in servers:
+    for thread, server in servers:
         server.shutdown()
-        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
+        thread.join(timeout=2)
 
 
 @pytest.fixture(autouse=True)
-def clean_sys_argv():
+def clean_sys_argv(monkeypatch):
     """Ensure sys.argv is clean for all tests."""
-    original = sys.argv[:]
-    sys.argv = ["test_script.py"]
-    yield
-    sys.argv = original
+    monkeypatch.setattr(sys, "argv", ["test_script.py"])
 
 
 def make_mock_connection(src_port, dest_port, src_user="src_user", dest_user="dest_user"):
@@ -150,7 +111,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 or "")
+        c.login(user, pwd)
         return c
 
     return mock_conn
@@ -163,31 +124,7 @@ def make_single_mock_connection(port):
 
     def mock_conn(host, user, pwd, oauth2_token=None):
         c = imaplib.IMAP4("localhost", port)
-        c.login(user, pwd or "")
+        c.login(user, pwd)
         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

@@ -1,341 +0,0 @@
-"""
-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

@@ -1,96 +0,0 @@
-"""
-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

@@ -1,286 +0,0 @@
-"""
-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

@@ -0,0 +1,848 @@
+"""
+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

@@ -0,0 +1,425 @@
+"""
+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

@@ -0,0 +1,256 @@
+"""
+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

@@ -1,565 +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 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_common.py

@@ -14,39 +14,41 @@ Tests cover:
 
 import os
 import sys
-from unittest.mock import Mock, patch
+from unittest.mock import Mock
 
-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__), "../../src")))
-
-from conftest import temp_env
-from utils import imap_common
+import imap_common
 
 
 class TestVerifyEnvVars:
     """Tests for verify_env_vars function."""
 
-    def test_all_vars_present(self):
+    def test_all_vars_present(self, monkeypatch):
         """Test returns True when all variables are set."""
-        with temp_env({"VAR1": "value1", "VAR2": "value2"}):
-            result = imap_common.verify_env_vars(["VAR1", "VAR2"])
-            assert result is True
+        monkeypatch.setenv("VAR1", "value1")
+        monkeypatch.setenv("VAR2", "value2")
 
-    def test_missing_vars(self, capsys):
+        result = imap_common.verify_env_vars(["VAR1", "VAR2"])
+        assert result is True
+
+    def test_missing_vars(self, monkeypatch, capsys):
         """Test returns False and prints error when variables are missing."""
-        with temp_env({}):
-            result = imap_common.verify_env_vars(["MISSING_VAR"])
-            assert result is False
+        monkeypatch.delenv("MISSING_VAR", raising=False)
 
-            captured = capsys.readouterr()
-            assert "MISSING_VAR" in captured.err
+        result = imap_common.verify_env_vars(["MISSING_VAR"])
+        assert result is False
 
-    def test_partial_vars_present(self, capsys):
+        captured = capsys.readouterr()
+        assert "MISSING_VAR" in captured.err
+
+    def test_partial_vars_present(self, monkeypatch, capsys):
         """Test with some variables present and some missing."""
-        with temp_env({"PRESENT": "value"}):
-            result = imap_common.verify_env_vars(["PRESENT", "MISSING"])
-            assert result is False
+        monkeypatch.setenv("PRESENT", "value")
+        monkeypatch.delenv("MISSING", raising=False)
+
+        result = imap_common.verify_env_vars(["PRESENT", "MISSING"])
+        assert result is False
 
 
 class TestGetImapConnection:
@@ -72,60 +74,6 @@ 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."""
@@ -320,6 +268,48 @@ 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."""
 
@@ -604,21 +594,20 @@ class TestAppendEmail:
 
         assert result is False
 
-    def test_append_exception_propagates(self):
-        """Test append propagates exceptions so callers can handle/log."""
+    def test_append_exception_returns_false(self):
+        """Test append returns False when exception occurs."""
         mock_conn = Mock()
         mock_conn.append.side_effect = Exception("Connection error")
 
-        import pytest
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            ensure_folder=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,
-            )
+        assert result is False
 
     def test_append_with_multiple_flags(self):
         """Test append with multiple flags."""
@@ -663,22 +652,14 @@ class TestGetMessageIdsInFolder:
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
-    def test_search_non_auth_exception_returns_empty(self):
-        """Test returns empty dict when search raises non-auth exception."""
+    def test_search_exception(self):
+        """Test returns empty dict when search raises exception."""
         mock_conn = Mock()
         mock_conn.uid.side_effect = Exception("Connection error")
 
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
-    def test_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()
@@ -725,8 +706,8 @@ class TestGetMessageIdsInFolder:
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
-    def test_fetch_non_auth_exception_continues(self):
-        """Test continues when fetch raises non-auth exception for a batch."""
+    def test_fetch_exception_for_batch(self):
+        """Test continues when fetch raises exception for a batch."""
         mock_conn = Mock()
         mock_conn.uid.side_effect = [
             ("OK", [b"1"]),
@@ -736,17 +717,6 @@ 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()
@@ -765,77 +735,3 @@ 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/test_imap_compare.py

@@ -1,257 +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
-
-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

@@ -1,302 +0,0 @@
-"""
-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_oauth2.py

@@ -0,0 +1,489 @@
+"""
+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_retry.py

@@ -1,188 +0,0 @@
-"""
-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_imap_emails.py

@@ -19,42 +19,38 @@ 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 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",
-    }
+import migrate_imap_emails
+from conftest import make_mock_connection
 
 
 class TestBasicMigration:
     """Tests for basic migration functionality."""
 
-    def test_single_email_migration(self, mock_server_factory):
+    def test_single_email_migration(self, mock_server_factory, monkeypatch):
         """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 = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        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):
+    def test_multiple_emails_migration(self, mock_server_factory, monkeypatch):
         """Test migrating multiple emails."""
         src_data = {
             "INBOX": [
@@ -67,9 +63,19 @@ class TestBasicMigration:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        migrate_imap_emails.main()
 
         assert len(dest_server.folders["INBOX"]) == 3
 
@@ -77,7 +83,7 @@ class TestBasicMigration:
 class TestDuplicateHandling:
     """Tests for duplicate detection and skipping."""
 
-    def test_skip_duplicate_by_message_id(self, mock_server_factory):
+    def test_skip_duplicate_by_message_id(self, mock_server_factory, monkeypatch):
         """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]}
@@ -85,14 +91,24 @@ class TestDuplicateHandling:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        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):
+    def test_migrate_non_duplicate(self, mock_server_factory, monkeypatch):
         """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"
@@ -102,9 +118,19 @@ class TestDuplicateHandling:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        migrate_imap_emails.main()
 
         # Should now have 2 messages
         assert len(dest_server.folders["INBOX"]) == 2
@@ -113,31 +139,52 @@ class TestDuplicateHandling:
 class TestDeleteFromSource:
     """Tests for delete-after-migration functionality."""
 
-    def test_delete_after_migration(self, mock_server_factory):
+    def test_delete_after_migration(self, mock_server_factory, monkeypatch):
         """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 = _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()
+        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()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert len(src_server.folders["INBOX"]) == 0
 
-    def test_no_delete_when_disabled(self, mock_server_factory):
+    def test_no_delete_when_disabled(self, mock_server_factory, monkeypatch):
         """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 = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
-            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",
+            # 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()
 
         assert len(dest_server.folders["INBOX"]) == 1
         assert len(src_server.folders["INBOX"]) == 1
@@ -146,21 +193,31 @@ class TestDeleteFromSource:
 class TestFolderHandling:
     """Tests for folder creation and multi-folder migration."""
 
-    def test_folder_creation(self, mock_server_factory):
+    def test_folder_creation(self, mock_server_factory, monkeypatch):
         """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 = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        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):
+    def test_multiple_folders_migration(self, mock_server_factory, monkeypatch):
         """Test migrating emails from multiple folders."""
         src_data = {
             "INBOX": [b"Subject: Inbox\r\nMessage-ID: <inbox>\r\n\r\nC"],
@@ -171,28 +228,49 @@ class TestFolderHandling:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        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):
+    def test_empty_folder_handling(self, mock_server_factory, monkeypatch):
         """Test that empty folders are handled gracefully."""
         src_data = {"INBOX": [], "Empty": []}
         dest_data = {"INBOX": []}
 
         _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py"]):
-            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)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        # Should complete without error
+        migrate_imap_emails.main()
 
 
 class TestPreserveFlags:
-    def test_preserve_flags_applied_on_copy(self, mock_server_factory):
+    def test_preserve_flags_applied_on_copy(self, mock_server_factory, monkeypatch):
         msg = b"Subject: Flagged\r\nMessage-ID: <f1@test>\r\n\r\nBody"
         src_data = {"INBOX": [msg]}
         dest_data = {"INBOX": []}
@@ -200,15 +278,25 @@ class TestPreserveFlags:
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
         src_server.folders["INBOX"][0]["flags"].add("\\Seen")
 
-        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()
+        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()
 
         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):
+    def test_preserve_flags_syncs_on_duplicate(self, mock_server_factory, monkeypatch):
         msg = b"Subject: DupFlags\r\nMessage-ID: <df1@test>\r\n\r\nBody"
         src_data = {"INBOX": [msg]}
         dest_data = {"INBOX": [msg]}
@@ -216,17 +304,27 @@ class TestPreserveFlags:
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
         src_server.folders["INBOX"][0]["flags"].add("\\Seen")
 
-        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()
+        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()
 
         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):
+    def test_gmail_mode_migrates_all_mail_and_applies_labels(self, mock_server_factory, monkeypatch):
         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"
 
@@ -239,108 +337,54 @@ class TestGmailModeLabels:
 
         _, 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()
+        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()
 
         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, capsys):
+    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",
         }
-        with temp_env(env):
-            with pytest.raises(SystemExit) as exc_info:
-                migrate_imap_emails.main()
+        monkeypatch.setattr(os, "environ", env)
+
+        with pytest.raises(SystemExit) as exc_info:
+            migrate_imap_emails.main()
 
         assert exc_info.value.code == 2
 
-    def test_missing_dest_credentials(self, capsys):
+    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",
         }
-        with temp_env(env):
-            with pytest.raises(SystemExit) as exc_info:
-                migrate_imap_emails.main()
+        monkeypatch.setattr(os, "environ", env)
+
+        with pytest.raises(SystemExit) as exc_info:
+            migrate_imap_emails.main()
 
         assert exc_info.value.code == 2
 
@@ -348,32 +392,38 @@ class TestConfigValidation:
 class TestMigrateErrorHandling:
     """Tests for error handling during migration."""
 
-    def test_connection_error_in_worker(self, mock_server_factory):
+    def test_connection_error_in_worker(self, mock_server_factory, monkeypatch):
         """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)
 
-        from unittest.mock import patch
+        # Proper setup for main thread connection
+        real_mock_conn = make_mock_connection(p1, p2)
 
-        original_get_imap_connection = imap_common.get_imap_connection
-
-        def side_effect(host, user, pwd, oauth2_token=None):
+        def side_effect(*args, **kwargs):
             if threading.current_thread() is threading.main_thread():
-                return original_get_imap_connection(host, user, pwd, oauth2_token)
+                return real_mock_conn(*args, **kwargs)
             return None  # Simulate connection failure in worker
 
-        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()
+        monkeypatch.setattr("imap_common.get_imap_connection", side_effect)
 
-    def test_select_error_in_worker(self, mock_server_factory):
+        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):
         """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": []}
@@ -388,44 +438,23 @@ class TestMigrateErrorHandling:
                 raise RuntimeError("Select failed")
             return original_select(self, mailbox, readonly)
 
-        from unittest.mock import patch
+        monkeypatch.setattr(imaplib.IMAP4, "select", side_effect_select)
 
-        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()
+        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))
 
-    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": []}
+        migrate_imap_emails.main()
 
-        _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):
+    def test_fetch_error_in_worker(self, mock_server_factory, monkeypatch):
         """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": []}
@@ -440,22 +469,28 @@ class TestMigrateErrorHandling:
                 raise RuntimeError("Fetch failed")
             return original_uid(self, command, *args)
 
-        from unittest.mock import patch
+        monkeypatch.setattr(imaplib.IMAP4, "uid", side_effect_uid)
 
-        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()
+        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()
 
         # Dest should be empty as fetch failed
         assert len(dest_server.folders["INBOX"]) == 0
 
-    def test_main_connection_failure(self):
+    def test_main_connection_failure(self, monkeypatch):
         """Test that main exits if initial connection fails."""
-        from unittest.mock import patch
+        monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: None)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -465,60 +500,45 @@ class TestMigrateErrorHandling:
             "DEST_IMAP_USERNAME": "dest_user",
             "DEST_IMAP_PASSWORD": "p",
         }
-        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
+        monkeypatch.setattr(os, "environ", env)
 
-    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"]),
-        ):
+        with pytest.raises(SystemExit) as exc:
             migrate_imap_emails.main()
-
-        captured = capsys.readouterr()
-        assert "Warning: Failed to load progress cache" in captured.out
+        assert exc.value.code == 1
 
 
 class TestTrashHandling:
     """Tests for trash folder related logic."""
 
-    def test_circular_trash_migration_prevention(self, mock_server_factory):
+    def test_circular_trash_migration_prevention(self, mock_server_factory, monkeypatch):
         """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)
 
-        from unittest.mock import patch
+        # Mock trash detection
+        monkeypatch.setattr("imap_common.detect_trash_folder", lambda conn: "Trash")
 
-        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()
+        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()
 
         # 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):
+    def test_deleted_moved_to_trash(self, mock_server_factory, monkeypatch):
         """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": []}
@@ -526,16 +546,22 @@ class TestTrashHandling:
         src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
         src_server.folders["Trash"] = []  # Create trash on source
 
-        from unittest.mock import patch
+        monkeypatch.setattr("imap_common.detect_trash_folder", lambda conn: "Trash")
 
-        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()
+        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()
 
         # Dest should have it
         assert len(dest_server.folders["INBOX"]) == 1
@@ -545,58 +571,6 @@ 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."""
 
@@ -675,7 +649,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):
+    def test_delete_orphan_emails_removes_extra_dest_emails(self, mock_server_factory, monkeypatch):
         """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"]}
@@ -690,17 +664,24 @@ class TestDestDeleteFunctionality:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        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()
+        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()
 
         # 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):
+    def test_dest_delete_disabled_keeps_extra_emails(self, mock_server_factory, monkeypatch):
         """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 = {
@@ -711,14 +692,22 @@ class TestDestDeleteFunctionality:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        env = _mock_migrate_env(p1, p2)
-        with temp_env(env), temp_argv(["migrate_imap_emails.py", "INBOX"]):
-            migrate_imap_emails.main()
+        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()
 
         # 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):
+    def test_dest_delete_empty_source_deletes_all(self, mock_server_factory, monkeypatch):
         """Test that if source folder is empty, all dest emails are deleted."""
         src_data = {"INBOX": []}
         dest_data = {
@@ -730,46 +719,17 @@ class TestDestDeleteFunctionality:
 
         _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
 
-        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()
+        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()
 
         # 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_migrate_resume.py

@@ -1,165 +0,0 @@
-"""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

@@ -1,248 +0,0 @@
-"""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/test_restore_imap_emails.py

@@ -9,28 +9,19 @@ 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_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",
-    }
+import imap_common
+import restore_cache
+import restore_imap_emails
+from conftest import make_single_mock_connection
 
 
 class TestLoadLabelsManifest:
@@ -45,7 +36,7 @@ class TestLoadLabelsManifest:
         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")
+        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
         assert result == manifest_data
 
     def test_load_manifest_new_format(self, tmp_path):
@@ -57,14 +48,14 @@ class TestLoadLabelsManifest:
         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")
+        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
         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 = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
+        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
         assert result == {}
 
     def test_load_invalid_json_manifest(self, tmp_path):
@@ -72,7 +63,7 @@ class TestLoadLabelsManifest:
         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")
+        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
         assert result == {}
 
 
@@ -156,15 +147,6 @@ 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."""
@@ -193,48 +175,168 @@ 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, capsys, tmp_path):
+    def test_missing_credentials(self, monkeypatch, capsys, tmp_path):
         """Test that missing credentials cause exit."""
-        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()
+        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()
 
         assert exc_info.value.code == 2
 
-    def test_missing_src_path(self, capsys):
+    def test_missing_src_path(self, monkeypatch, capsys):
         """Test that missing source path causes exit."""
         env = {
             "DEST_IMAP_HOST": "localhost",
             "DEST_IMAP_USERNAME": "user",
             "DEST_IMAP_PASSWORD": "pass",
         }
-        with temp_env(env), temp_argv(["restore_imap_emails.py"]):
-            with pytest.raises(SystemExit) as exc_info:
-                restore_imap_emails.main()
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr(sys, "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, capsys):
+    def test_nonexistent_src_path(self, monkeypatch, capsys):
         """Test that non-existent source path causes exit."""
         env = {
             "DEST_IMAP_HOST": "localhost",
             "DEST_IMAP_USERNAME": "user",
             "DEST_IMAP_PASSWORD": "pass",
         }
-        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()
+        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()
 
         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, tmp_path):
+    def test_restore_single_folder(self, single_mock_server, monkeypatch, tmp_path):
         """Test restoring a single folder."""
         # Create backup structure
         inbox = tmp_path / "INBOX"
@@ -254,39 +356,23 @@ Body content.
         src_data = {"INBOX": []}
         server, port = single_mock_server(src_data)
 
-        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()
+        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))
 
-    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")
+        # Run restore
+        restore_imap_emails.main()
 
-        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):
+    def test_restore_single_folder_full_restore_flag(self, single_mock_server, monkeypatch, tmp_path):
         """Smoke test: --full-restore flag is accepted."""
         inbox = tmp_path / "INBOX"
         inbox.mkdir()
@@ -304,12 +390,20 @@ Body content.
         src_data = {"INBOX": []}
         _server, port = single_mock_server(src_data)
 
-        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()
+        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()
 
     def test_restore_with_labels_manifest(self, tmp_path):
         """Test that labels manifest is loaded correctly."""
@@ -322,49 +416,48 @@ Body content.
         manifest_path.write_text(json.dumps(manifest_data))
 
         # Load and verify
-        result = imap_common.load_manifest(str(tmp_path), "labels_manifest.json")
+        result = restore_imap_emails.load_labels_manifest(str(tmp_path))
         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")
 
-        dest_data = {"INBOX": []}
-        server, port = single_mock_server(dest_data)
+class TestEmailExistsInFolder:
+    """Tests for duplicate detection."""
 
-        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()
+    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)
 
-        assert imap_common.FOLDER_RESTORED_UNLABELED in server.folders
-        assert len(server.folders[imap_common.FOLDER_RESTORED_UNLABELED]) == 1
+        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
+        assert result is True
 
-    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")
+    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)
 
-        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)
+        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
+        assert result is False
 
-        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()
+    def test_email_exists_no_message_id(self, monkeypatch):
+        """Test with no message ID."""
+        mock_conn = MagicMock()
 
-        assert len(server.folders["INBOX"]) == 1
-        assert b"Message-ID: <keep@test>" in server.folders["INBOX"][0]["content"]
+        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
 
 
 class TestRestoreProgressCache:
@@ -407,60 +500,6 @@ 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."""
 
@@ -493,48 +532,107 @@ class TestGetLabelsFromManifest:
         assert result == []
 
 
-class TestRestoreE2EHelpers:
-    """End-to-end tests for restore helper functions using the mock IMAP server."""
+class TestLabelToFolder:
+    """Tests for label_to_folder function."""
 
-    def test_upload_email_success(self, single_mock_server):
-        dest_data = {"INBOX": []}
-        server, port = single_mock_server(dest_data)
+    def test_inbox_label(self):
+        """Test INBOX label conversion."""
+        result = restore_imap_emails.label_to_folder("INBOX")
+        assert result == "INBOX"
 
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
+    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"
 
-        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"',
+    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,
         )
 
-        assert result == restore_imap_emails.UploadResult.SUCCESS
-        assert len(server.folders["INBOX"]) == 1
-        conn.logout()
+        assert captured["folder_name"] == imap_common.FOLDER_RESTORED_UNLABELED
+        assert captured["folder_name"] != "[Gmail]/Drafts"
 
-    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)
 
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
+class TestSyncFlagsOnExisting:
+    """Tests for sync_flags_on_existing function."""
 
-        imap_common.sync_flags_on_existing(conn, "INBOX", "<flag@test>", "\\Seen \\Flagged", 1000)
+    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)
 
-        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()
+        # 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()
 
 
 class TestDestDeleteRestoreArgument:
@@ -585,7 +683,7 @@ class TestDestDeleteRestoreFunctionality:
         # Local backup only has one email
         local_msg_ids = {"<keep@test>"}
 
-        deleted = imap_common.delete_orphan_emails(conn, "INBOX", local_msg_ids)
+        deleted = restore_imap_emails.delete_orphan_emails_from_dest(conn, "INBOX", local_msg_ids)
 
         assert deleted == 2
         # Verify only 1 email remains
@@ -613,7 +711,7 @@ class TestDestDeleteRestoreFunctionality:
         # Empty local backup
         local_msg_ids = set()
 
-        deleted = imap_common.delete_orphan_emails(conn, "INBOX", local_msg_ids)
+        deleted = restore_imap_emails.delete_orphan_emails_from_dest(conn, "INBOX", local_msg_ids)
 
         assert deleted == 2
         assert len(server.folders["INBOX"]) == 0
@@ -639,7 +737,7 @@ class TestDestDeleteRestoreFunctionality:
         # Local has both emails
         local_msg_ids = {"<e1@test>", "<e2@test>"}
 
-        deleted = imap_common.delete_orphan_emails(conn, "INBOX", local_msg_ids)
+        deleted = restore_imap_emails.delete_orphan_emails_from_dest(conn, "INBOX", local_msg_ids)
 
         assert deleted == 0
         assert len(server.folders["INBOX"]) == 2
@@ -686,7 +784,9 @@ 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, tmp_path):
+    def test_dest_delete_enabled_via_env_var_main_deletes_all_in_folder(
+        self, single_mock_server, monkeypatch, tmp_path
+    ):
         """End-to-end: DEST_DELETE=true triggers deletion when local folder has no .eml files."""
         dest_data = {
             "INBOX": [
@@ -700,74 +800,17 @@ class TestDestDeleteRestoreFunctionality:
         backup_root = tmp_path / "backup"
         (backup_root / "INBOX").mkdir(parents=True)
 
-        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()
+        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()
 
         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/utils/test_imap_compress.py

@@ -1,394 +0,0 @@
-"""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,3 +1,4 @@
+import re
 import socketserver
 import threading
 
@@ -14,15 +15,7 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
     def handle(self):
         self.wfile.write(b"* OK [CAPABILITY IMAP4rev1] Mock IMAP Server Ready\r\n")
         self.selected_folder = None
-
-        # 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
+        self.current_folders = self.server.folders
 
         while True:
             try:
@@ -38,73 +31,9 @@ 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
@@ -113,16 +42,6 @@ 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())
@@ -130,40 +49,23 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
 
                 elif cmd == "SELECT":
                     folder = args.strip().strip('"')
-
-                    # 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:
+                    if folder in self.current_folders:
                         self.selected_folder = folder
-                        count = len(self.current_folders.get(folder, []))  # Default empty for retry folders
+                        count = len(self.current_folders[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")
                         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('"')
-
-                    is_retry_folder = "RETRY_" in folder
-
-                    if folder in self.current_folders or is_retry_folder:
+                    if folder in self.current_folders:
                         self.selected_folder = folder
-                        count = len(self.current_folders.get(folder, []))
+                        count = len(self.current_folders[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")
@@ -178,42 +80,6 @@ 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]
@@ -259,10 +125,7 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                             valid_uids.append(str(m["uid"]))
 
                         uids_str = " ".join(valid_uids)
-                        if uids_str:
-                            self.wfile.write(f"* SEARCH {uids_str}\r\n".encode())
-                        else:
-                            self.wfile.write(b"* SEARCH\r\n")
+                        self.wfile.write(f"* SEARCH {uids_str}\r\n".encode())
                         self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
 
                     elif sub_cmd == "STORE":
@@ -429,6 +292,50 @@ 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:
@@ -504,11 +411,7 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 else:
                     self.send_response(tag, "BAD Command not recognized")
 
-            except Exception as e:
-                print(f"MockServer EXCEPTION: {e}")
-                import traceback
-
-                traceback.print_exc()
+            except Exception:
                 break
 
     def send_response(self, tag, message):
@@ -534,19 +437,9 @@ class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
             self.folders = {"INBOX": []}
 
 
-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
-
+def start_server_thread(port=10143, initial_folders=None):
     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 server, actual_port
+    return t, server

tools/mock_oauth_server.py

@@ -1,59 +0,0 @@
-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