111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
"""In-process SMTP server for E2E email assertions."""
|
|
|
|
import asyncio
|
|
import socket
|
|
import threading
|
|
import time
|
|
from email import message_from_bytes, policy
|
|
from email.message import EmailMessage, Message
|
|
|
|
from aiosmtpd.controller import Controller
|
|
|
|
|
|
class _CaptureHandler:
|
|
"""
|
|
aiosmtpd handler that parses every incoming message into an ``EmailMessage``.
|
|
|
|
The default policy yields a ``Message`` instance, which loses structure
|
|
when wrapped into a new ``EmailMessage``. Parsing with ``policy.default``
|
|
keeps multipart payloads intact so tests can introspect the HTML body
|
|
via ``walk()``/``get_payload(decode=True)``.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.messages: list[EmailMessage | Message] = []
|
|
|
|
async def handle_DATA(self, server, session, envelope):
|
|
message = message_from_bytes(envelope.content, policy=policy.default)
|
|
self.messages.append(message)
|
|
return '250 OK'
|
|
|
|
|
|
def _reserve_port(host: str = '127.0.0.1') -> int:
|
|
"""Reserve a free TCP port on the given host."""
|
|
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
probe.bind((host, 0))
|
|
port = probe.getsockname()[1]
|
|
probe.close()
|
|
return port
|
|
|
|
|
|
class SmtpTestServer:
|
|
"""
|
|
Wraps aiosmtpd Controller with a dedicated asyncio loop thread for pytest compatibility.
|
|
|
|
Attributes:
|
|
host: Bind host (127.0.0.1).
|
|
port: Listening port after start().
|
|
messages: Captured outbound messages.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.host = '127.0.0.1'
|
|
self.port: int | None = None
|
|
self._handler = _CaptureHandler()
|
|
self.messages: list[EmailMessage | Message] = self._handler.messages
|
|
self._controller: Controller | None = None
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
self._thread: threading.Thread | None = None
|
|
|
|
def start(self) -> None:
|
|
"""Start the SMTP controller on a reserved port in a background event loop."""
|
|
self.port = _reserve_port(self.host)
|
|
self._loop = asyncio.new_event_loop()
|
|
self._controller = Controller(
|
|
self._handler,
|
|
hostname=self.host,
|
|
port=self.port,
|
|
loop=self._loop,
|
|
ready_timeout=30,
|
|
)
|
|
|
|
def _run():
|
|
asyncio.set_event_loop(self._loop)
|
|
if self._controller is None:
|
|
raise RuntimeError('SMTP controller not initialized')
|
|
self._controller.start()
|
|
|
|
self._thread = threading.Thread(target=_run, name='e2e-smtp', daemon=True)
|
|
self._thread.start()
|
|
self._wait_until_ready()
|
|
|
|
def _wait_until_ready(self, timeout: float = 10.0) -> None:
|
|
"""Poll until the SMTP listener accepts TCP connections."""
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if self.port is None:
|
|
time.sleep(0.05)
|
|
continue
|
|
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
try:
|
|
if probe.connect_ex((self.host, self.port)) == 0:
|
|
return
|
|
finally:
|
|
probe.close()
|
|
time.sleep(0.05)
|
|
raise TimeoutError('SMTP test server did not become ready')
|
|
|
|
def stop(self) -> None:
|
|
"""Stop the controller and background event loop."""
|
|
if self._controller is not None:
|
|
self._controller.stop()
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=5)
|
|
self._controller = None
|
|
self._loop = None
|
|
self._thread = None
|
|
|
|
def clear(self) -> None:
|
|
"""Remove all captured messages."""
|
|
self.messages.clear()
|