From c6a568857f729ec190f5aefa22a03cbbb27d080e Mon Sep 17 00:00:00 2001 From: Javier Callico Date: Tue, 10 Feb 2026 11:06:18 -0500 Subject: [PATCH] Add connection proxy wrapper for IMAP retry functionality (#34) --- src/imap_common.py | 3 +- src/imap_retry.py | 80 ++++++++++++++ test/conftest.py | 32 ++++-- test/test_count_imap_emails.py | 2 - test/test_imap_retry.py | 188 +++++++++++++++++++++++++++++++++ tools/mock_imap_server.py | 124 ++++++++++++++++++++-- 6 files changed, 407 insertions(+), 22 deletions(-) create mode 100644 src/imap_retry.py create mode 100644 test/test_imap_retry.py diff --git a/src/imap_common.py b/src/imap_common.py index 9006653..5e0a102 100644 --- a/src/imap_common.py +++ b/src/imap_common.py @@ -18,6 +18,7 @@ from email.header import decode_header from email.parser import BytesParser import imap_oauth2 +import imap_retry import restore_cache # Standard IMAP flags @@ -289,7 +290,7 @@ def get_imap_connection(host, user, password=None, oauth2_token=None): conn.authenticate("XOAUTH2", lambda _: auth_string.encode()) else: conn.login(user, password) - return conn + return imap_retry.ConnectionProxy(conn, log_fn=safe_print) except Exception as e: print(f"Connection error to {host}: {e}") return None diff --git a/src/imap_retry.py b/src/imap_retry.py new file mode 100644 index 0000000..81aa57e --- /dev/null +++ b/src/imap_retry.py @@ -0,0 +1,80 @@ +""" +IMAP Retry Logic + +Transparent retry wrapper for IMAP connections that handles transient +server errors (e.g. Microsoft 365 "Server Busy") with exponential backoff. +""" + +from __future__ import annotations + +import time + + +class ConnectionProxy: + """Transparent proxy that retries IMAP commands on transient server errors. + + Wraps an imaplib.IMAP4 or IMAP4_SSL connection. For methods in + RETRYABLE_METHODS that return (typ, data) tuples, retries on transient + errors with exponential backoff. + """ + + TRANSIENT_PATTERNS = [b"UNAVAILABLE", b"Server Busy", b"try again", b"THROTTLED"] + + # Methods that are safe to retry and return (typ, data) + RETRYABLE_METHODS = frozenset( + { + "uid", + "select", + "search", + "fetch", + "append", + "store", + "list", + "create", + "expunge", + "noop", + } + ) + + def __init__(self, conn, max_retries=3, initial_wait=5, log_fn=print): + if max_retries < 1: + raise ValueError(f"max_retries must be >= 1, got {max_retries}") + if initial_wait < 0: + raise ValueError(f"initial_wait must be >= 0, got {initial_wait}") + self._conn = conn + self._max_retries = max_retries + self._initial_wait = initial_wait + self._log_fn = log_fn + + @classmethod + def _is_transient_error(cls, data): + """Check if IMAP response data contains transient error patterns.""" + for item in data: + if isinstance(item, bytes): + for pattern in cls.TRANSIENT_PATTERNS: + if pattern in item: + return True + return False + + def __getattr__(self, name): + attr = getattr(self._conn, name) + if name not in self.RETRYABLE_METHODS or not callable(attr): + return attr + + def wrapper(*args, **kwargs): + last_result = None + for attempt in range(self._max_retries): + result = attr(*args, **kwargs) + if not isinstance(result, tuple) or len(result) < 2: + return result + typ, data = result[0], result[1] + if typ == "OK" or not self._is_transient_error(data): + return result + last_result = result + if attempt + 1 < self._max_retries: + wait = self._initial_wait * (2**attempt) # 5s, 10s, 20s + self._log_fn(f"Server busy, retrying in {wait}s... (attempt {attempt + 1}/{self._max_retries})") + time.sleep(wait) + return last_result + + return wrapper diff --git a/test/conftest.py b/test/conftest.py index 3baac9e..59b6c99 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -33,11 +33,19 @@ def create_server_pair(src_data=None, dest_data=None): while p2 == p1: p2 = get_free_port() - src_t, src_s = start_server_thread(p1, src_data) - dest_t, dest_s = start_server_thread(p2, dest_data) + src_s, p1_actual = start_server_thread(p1, src_data) + dest_s, p2_actual = start_server_thread(p2, dest_data) time.sleep(0.3) - return (src_t, src_s, p1), (dest_t, dest_s, p2) + # We don't have direct access to thread object anymore as it is managed internally by TCPServer/ThreadingMixin or start_server_thread wrapper? + # Actually start_server_thread in tools/mock_imap_server.py now returns (server, port). + # The server object is a ThreadingMixIn so it handles per-request threads, but the main serve_forever loop is in a thread we created. + # But wait, my modified start_server_thread does: t.start(), return server, actual_port. + # It dropped returning 't'. I need to fix that or update consumers to not need 't'. + # Shutdown() stops serve_forever loop. join() is nice but maybe not strictly required if we trust shutdown. + # However, to avoid ResourceWarnings, we might want 't'. + + return (None, src_s, p1_actual), (None, dest_s, p2_actual) def shutdown_server_pair(src_tuple, dest_tuple): @@ -45,9 +53,13 @@ def shutdown_server_pair(src_tuple, dest_tuple): src_t, src_s, _ = src_tuple dest_t, dest_s, _ = dest_tuple src_s.shutdown() + src_s.server_close() # Explicitly close socket dest_s.shutdown() - src_t.join(timeout=2) - dest_t.join(timeout=2) + dest_s.server_close() + if src_t: + src_t.join(timeout=2) + if dest_t: + dest_t.join(timeout=2) @pytest.fixture @@ -82,16 +94,16 @@ def single_mock_server(): def _create(initial_data=None): port = get_free_port() - thread, server = start_server_thread(port, initial_data) + server, actual_port = start_server_thread(port, initial_data) time.sleep(0.3) - servers.append((thread, server)) - return server, port + servers.append(server) + return server, actual_port yield _create - for thread, server in servers: + for server in servers: server.shutdown() - thread.join(timeout=2) + server.server_close() @contextmanager diff --git a/test/test_count_imap_emails.py b/test/test_count_imap_emails.py index 3814ba7..16b743b 100644 --- a/test/test_count_imap_emails.py +++ b/test/test_count_imap_emails.py @@ -196,7 +196,6 @@ class TestImapCommonHelpers: with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP): conn = imap_common.get_imap_connection("host", "user", oauth2_token="token") - assert isinstance(conn, FakeIMAP) assert conn.auth_called is True assert conn.login_called is False @@ -217,7 +216,6 @@ class TestImapCommonHelpers: with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP): conn = imap_common.get_imap_connection("host", "user", password="pass") - assert isinstance(conn, FakeIMAP) assert conn.login_called is True assert conn.auth_called is False diff --git a/test/test_imap_retry.py b/test/test_imap_retry.py new file mode 100644 index 0000000..0993a12 --- /dev/null +++ b/test/test_imap_retry.py @@ -0,0 +1,188 @@ +""" +Tests for imap_retry.py + +Tests cover: +- Transient error detection +- ConnectionProxy transparent proxying +- Retry with exponential backoff on transient errors +- Pass-through for non-retryable methods +- Pass-through for non-transient errors +""" + +import imaplib +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tools"))) + +import imap_retry +from mock_imap_server import start_server_thread as start_mock_server + + +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) diff --git a/tools/mock_imap_server.py b/tools/mock_imap_server.py index 5743f43..609f793 100644 --- a/tools/mock_imap_server.py +++ b/tools/mock_imap_server.py @@ -1,4 +1,3 @@ -import re import socketserver import threading @@ -15,7 +14,15 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): def handle(self): self.wfile.write(b"* OK [CAPABILITY IMAP4rev1] Mock IMAP Server Ready\r\n") self.selected_folder = None - self.current_folders = self.server.folders + + # Check if server has folders, if not, initialize + if hasattr(self.server, "folders"): + self.current_folders = self.server.folders + else: + # Just in case it's used standalone without the threaded wrapper + self.current_folders = {"INBOX": []} + + self.retry_state = {} # key -> count while True: try: @@ -31,9 +38,73 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): cmd = parts[1].upper() args = parts[2] if len(parts) > 2 else "" + # Check for RETRY_N injection in args (for testing imap_retry) + # Pattern: RETRY_N where N is integer + import re + + retry_match = re.search(r"RETRY_(\d+)", args) + if retry_match: + count = int(retry_match.group(1)) + # Create a unique key for this command invocation context + # Just using command + args is roughly sufficient + key = f"{cmd}_{args}" + current_fails = self.retry_state.get(key, 0) + + if current_fails < count: + # Logic Fix: Only increment if we are going to fail + self.retry_state[key] = current_fails + 1 + self.send_response(tag, "NO [UNAVAILABLE] Server Busy (Simulated)") + continue + # Else fall through to normal command processing + if cmd == "LOGIN": self.send_response(tag, "OK LOGIN completed") + elif cmd == "APPEND": + # Check for literal size + import re + + size_match = re.search(r"\{(\d+)\}$", args) + data = b"" + if size_match: + size = int(size_match.group(1)) + # Send continuation + self.wfile.write(b"+\r\n") + self.wfile.flush() + # Read data + data = self.rfile.read(size) + + # Extract mailbox name to store the message + # Simplistic parsing: check for quoted or unquoted first argument + mailbox = "INBOX" # Default fallback + clean_args = args.lstrip() + if clean_args.startswith('"'): + end_quote = clean_args.find('"', 1) + if end_quote != -1: + mailbox = clean_args[1:end_quote] + else: + mailbox = clean_args.split(" ")[0] + + if mailbox not in self.current_folders: + self.current_folders[mailbox] = [] + + # Parse flags if present: look for (...) + msg_flags = set() + flag_match = re.search(r"\(([^)]+)\)", args) + if flag_match: + # e.g. (\Seen \Deleted) + flag_content = flag_match.group(1) + msg_flags = set(flag_content.split()) + + # Store message + new_uid = len(self.current_folders[mailbox]) + 1 + msg_obj = {"uid": new_uid, "flags": msg_flags, "content": data} + self.current_folders[mailbox].append(msg_obj) + + # Always succeed for test if we passed retry logic + self.wfile.write(tag.encode() + b" OK [APPENDUID 1 100] APPEND completed\r\n") + self.wfile.flush() + elif cmd == "LOGOUT": self.send_response(tag, "OK LOGOUT completed") break @@ -46,8 +117,12 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): # 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()) @@ -55,23 +130,40 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): elif cmd == "SELECT": folder = args.strip().strip('"') - if folder in self.current_folders: + + # For retry testing: if folder starts with NO, fail immediately + if folder.startswith("NO"): + # e.g. "NO [UNAVAILABLE]" + self.send_response(tag, folder) + continue + + # Allow RETRY_N folders to succeed if they passed the retry intercept logic + is_retry_folder = "RETRY_" in folder + + if folder in self.current_folders or is_retry_folder: self.selected_folder = folder - count = len(self.current_folders[folder]) + count = len(self.current_folders.get(folder, [])) # Default empty for retry folders self.wfile.write(f"* {count} EXISTS\r\n".encode()) self.wfile.write(f"* {count} RECENT\r\n".encode()) self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n") self.wfile.write(b"* OK [UIDVALIDITY 1] UIDs valid\r\n") self.send_response(tag, "OK [READ-WRITE] SELECT completed") + elif folder == "EMPTY": + # Special case for "empty_data" test in retry + self.wfile.write(b"* 0 EXISTS\r\n") + self.send_response(tag, "OK SELECT completed") else: self.send_response(tag, "NO [NONEXISTENT] Folder not found") elif cmd == "EXAMINE": # EXAMINE is like SELECT but read-only folder = args.strip().strip('"') - if folder in self.current_folders: + + is_retry_folder = "RETRY_" in folder + + if folder in self.current_folders or is_retry_folder: self.selected_folder = folder - count = len(self.current_folders[folder]) + count = len(self.current_folders.get(folder, [])) self.wfile.write(f"* {count} EXISTS\r\n".encode()) self.wfile.write(f"* {count} RECENT\r\n".encode()) self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n") @@ -412,7 +504,11 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): else: self.send_response(tag, "BAD Command not recognized") - except Exception: + except Exception as e: + print(f"MockServer EXCEPTION: {e}") + import traceback + + traceback.print_exc() break def send_response(self, tag, message): @@ -438,9 +534,19 @@ class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer): self.folders = {"INBOX": []} -def start_server_thread(port=10143, initial_folders=None): +def start_server_thread(port=0, initial_folders=None): + # Use port 0 to let OS select free port + if isinstance(port, dict): + # Handle legacy call where port was inadvertently passed as dict + initial_folders = port + port = 0 + server = MockIMAPServer(("localhost", port), MockIMAPHandler, initial_folders) + + # Get the actual port if 0 was used + actual_port = server.server_address[1] + t = threading.Thread(target=server.serve_forever) t.daemon = True t.start() - return t, server + return server, actual_port