imap-migration-tools/test/conftest.py

"""
Shared pytest fixtures and utilities for IMAP migration tools tests.
"""

import imaplib
import os
import socket
import sys
import time
from contextlib import contextmanager

import pytest

# Ensure src/tools are in path
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
from mock_oauth_server import start_server_thread as start_oauth_server_thread


def get_free_port():
    """Get a free port on localhost."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", 0))
        return s.getsockname()[1]


def create_server_pair(src_data=None, dest_data=None):
    """Create a pair of mock IMAP servers with given initial data."""
    p1 = get_free_port()
    p2 = get_free_port()
    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)
    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)


def shutdown_server_pair(src_tuple, dest_tuple):
    """Shutdown a pair of mock servers."""
    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)


@pytest.fixture
def mock_server_factory():
    """
    Factory fixture that creates mock IMAP server pairs.
    Automatically cleans up all servers after the test.
    """
    servers = []

    def _create(src_data=None, dest_data=None):
        src_tuple, dest_tuple = create_server_pair(src_data, dest_data)
        servers.append((src_tuple, dest_tuple))
        src_server = src_tuple[1]
        dest_server = dest_tuple[1]
        src_port = src_tuple[2]
        dest_port = dest_tuple[2]
        return src_server, dest_server, src_port, dest_port

    yield _create

    for src_tuple, dest_tuple in servers:
        shutdown_server_pair(src_tuple, dest_tuple)


@pytest.fixture
def single_mock_server():
    """
    Creates a single mock IMAP server for testing scripts that only need one server.
    """
    servers = []

    def _create(initial_data=None):
        port = get_free_port()
        server, actual_port = start_server_thread(port, initial_data)
        time.sleep(0.3)
        servers.append(server)
        return server, actual_port

    yield _create

    for 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


@pytest.fixture(autouse=True)
def clean_sys_argv():
    """Ensure sys.argv is clean for all tests."""
    original = sys.argv[:]
    sys.argv = ["test_script.py"]
    yield
    sys.argv = original


def make_mock_connection(src_port, dest_port, src_user="src_user", dest_user="dest_user"):
    """
    Creates a mock connection function that routes to the correct server based on username.
    """

    def mock_conn(host, user, pwd, oauth2_token=None):
        if user == src_user:
            port = src_port
        elif user == dest_user:
            port = dest_port
        else:
            raise ValueError(f"Unknown user: {user}")
        c = imaplib.IMAP4("localhost", port)
        c.login(user, pwd or "")
        return c

    return mock_conn


def make_single_mock_connection(port):
    """
    Creates a mock connection function for a single server.
    """

    def mock_conn(host, user, pwd, oauth2_token=None):
        c = imaplib.IMAP4("localhost", port)
        c.login(user, pwd or "")
        return c

    return mock_conn


@pytest.fixture
def mock_oauth_server():
    """Starts a mock OAuth2 server for token and discovery endpoints."""
    thread, server = start_oauth_server_thread(0)
    host, port = server.server_address
    base_url = f"http://{host}:{port}"

    yield base_url

    server.shutdown()
    thread.join(timeout=2)


__all__ = [
    "mock_server_factory",
    "single_mock_server",
    "make_mock_connection",
    "make_single_mock_connection",
    "mock_oauth_server",
    "temp_env",
    "temp_argv",
]