imap-migration-tools/test/test_imap_session.py
Nathan Moinvaziri e932190618 Add imap_session module for centralized OAuth2 refresh (#21)
* Add .claude to gitignore

* Add imap_session module for centralized OAuth2 refresh

Create new imap_session.py module that combines imap_common and
imap_oauth2 for unified session management:

- ensure_connection(): Refreshes OAuth2 token and ensures connection health
- ensure_folder_session(): Ensures connection AND folder selection

Updated backup, migrate, and restore scripts to use the new module,
simplifying the proactive token refresh pattern in process_batch
functions.

* Move is_token_expired_error to imap_oauth2 module

Relocates the OAuth2 token expiration detection function to the oauth2-specific module where it logically belongs.

* Add retry logic for auth errors during email processing

Extract single-UID processing into dedicated functions with retry capability
for recoverable authentication failures. Broadens auth error detection beyond
token expiration to include session loss and authentication failures.

* Cache Microsoft tenant discovery results

Avoids repeated OpenID Connect discovery requests for the same email domain
during OAuth2 authentication flows. Normalizes domain to lowercase for
consistent cache keys.

* Add tests for is_token_expired_error and is_auth_error

Covers token expiration, session invalidation, and authentication failure
detection with case-insensitive matching verification. Ensures non-auth
errors like connection and timeout issues are correctly excluded.

* Add tests for imap_session module

Covers ensure_connection and ensure_folder_session functions including
OAuth2 token refresh, folder selection, and error handling scenarios.
Verifies connection change detection triggers folder re-selection.

* Add tests for auth error re-raising in get_message_ids_in_folder

Verifies that authentication errors are propagated to callers for
reconnection handling while non-auth errors return empty results.
Clarifies test names to distinguish error handling behaviors.
2026-02-06 13:26:44 -05:00

269 lines
11 KiB
Python

"""
Tests for imap_session.py
Tests cover:
- ensure_connection() behavior with and without OAuth2
- ensure_folder_session() connection change detection
- Folder selection and re-selection logic
- Error handling for connection and folder failures
"""
import os
import sys
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
import imap_session
class TestEnsureConnection:
"""Tests for ensure_connection function."""
def test_with_oauth2_calls_refresh(self):
"""Test that OAuth2 token refresh is called when oauth2 config is present."""
conf = {
"host": "imap.example.com",
"user": "user@example.com",
"password": "pass",
"oauth2_token": "old_token",
"oauth2": {
"provider": "microsoft",
"client_id": "client-id",
"email": "user@example.com",
},
}
mock_conn = MagicMock()
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
with patch.object(
imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn
) as mock_ensure:
result = imap_session.ensure_connection(mock_conn, conf)
mock_refresh.assert_called_once_with(conf, "old_token")
mock_ensure.assert_called_once_with(mock_conn, conf)
assert result is mock_conn
def test_without_oauth2_skips_refresh(self):
"""Test that token refresh is skipped when oauth2 config is not present."""
conf = {
"host": "imap.example.com",
"user": "user@example.com",
"password": "pass",
"oauth2": None,
}
mock_conn = MagicMock()
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
imap_session.ensure_connection(mock_conn, conf)
mock_refresh.assert_not_called()
def test_without_oauth2_key_skips_refresh(self):
"""Test that token refresh is skipped when oauth2 key is missing entirely."""
conf = {
"host": "imap.example.com",
"user": "user@example.com",
"password": "pass",
}
mock_conn = MagicMock()
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
imap_session.ensure_connection(mock_conn, conf)
mock_refresh.assert_not_called()
def test_returns_connection_from_ensure_connection_from_conf(self):
"""Test returns the connection from ensure_connection_from_conf."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn):
result = imap_session.ensure_connection(old_conn, conf)
assert result is new_conn
def test_returns_none_when_connection_fails(self):
"""Test returns None when ensure_connection_from_conf fails."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=None):
result = imap_session.ensure_connection(MagicMock(), conf)
assert result is None
def test_with_none_connection(self):
"""Test works correctly when passed None as connection."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
new_conn = MagicMock()
with patch.object(
imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn
) as mock_ensure:
result = imap_session.ensure_connection(None, conf)
mock_ensure.assert_called_once_with(None, conf)
assert result is new_conn
class TestEnsureFolderSession:
"""Tests for ensure_folder_session function."""
def test_same_connection_no_folder_select(self):
"""Test that folder is not re-selected when connection stays the same."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
mock_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=mock_conn):
result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
assert result_conn is mock_conn
assert success is True
mock_conn.select.assert_not_called()
def test_new_connection_selects_folder(self):
"""Test that folder is selected when connection changes."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
def test_new_connection_selects_folder_readwrite(self):
"""Test that folder is selected with readonly=False when specified."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(
old_conn, conf, "[Gmail]/All Mail", readonly=False
)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"[Gmail]/All Mail"', readonly=False)
def test_none_connection_returns_failure(self):
"""Test returns (None, False) when ensure_connection returns None."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
mock_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=None):
result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
assert result_conn is None
assert success is False
def test_folder_select_failure_returns_false(self):
"""Test returns (conn, False) when folder selection fails."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
new_conn.select.side_effect = Exception("Folder not found")
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "NonExistent", readonly=True)
assert result_conn is new_conn
assert success is False
def test_none_old_connection_selects_folder(self):
"""Test that folder is selected when old connection was None."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(None, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
def test_folder_name_with_spaces(self):
"""Test folder selection with folder names containing spaces."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "My Folder", readonly=True)
assert success is True
new_conn.select.assert_called_once_with('"My Folder"', readonly=True)
def test_folder_select_imap_error(self):
"""Test handles IMAP-specific errors during folder selection."""
conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
old_conn = MagicMock()
new_conn = MagicMock()
new_conn.select.side_effect = Exception("IMAP error: NO SELECT failed")
with patch.object(imap_session, "ensure_connection", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is False
class TestEnsureFolderSessionWithOAuth2:
"""Integration-style tests for ensure_folder_session with OAuth2."""
def test_oauth2_refresh_triggers_folder_reselect(self):
"""Test that OAuth2 token refresh (new connection) triggers folder reselection."""
conf = {
"host": "outlook.office365.com",
"user": "user@example.com",
"password": "oauth2_token_here",
"oauth2_token": "old_token",
"oauth2": {
"provider": "microsoft",
"client_id": "client-id",
"email": "user@example.com",
},
}
old_conn = MagicMock()
new_conn = MagicMock()
# Simulate token refresh causing a new connection
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"):
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn):
result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
assert result_conn is new_conn
assert success is True
new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
def test_healthy_connection_no_reselect(self):
"""Test that healthy connection (same object) doesn't reselect folder."""
conf = {
"host": "outlook.office365.com",
"user": "user@example.com",
"password": "oauth2_token_here",
"oauth2_token": "valid_token",
"oauth2": {
"provider": "microsoft",
"client_id": "client-id",
"email": "user@example.com",
},
}
mock_conn = MagicMock()
# Simulate healthy connection (same object returned)
with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"):
with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
assert result_conn is mock_conn
assert success is True
mock_conn.select.assert_not_called()