-
Notifications
You must be signed in to change notification settings - Fork 85
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
3 changed files
with
251 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,107 @@ | ||
from typing import Optional, Dict | ||
from enum import Enum | ||
import asyncio | ||
from datetime import datetime | ||
from uuid import uuid4 | ||
|
||
from .schemas import MessageSchema | ||
from .fastmail import FastMail | ||
|
||
class EmailStatus(Enum): | ||
QUEUED = "queued" | ||
SENDING = "sending" | ||
SENT = "sent" | ||
FAILED = "failed" | ||
|
||
class QueuedEmail: | ||
def __init__(self, message: MessageSchema, template_name: Optional[str] = None): | ||
self.id = str(uuid4()) | ||
self.message = message | ||
self.template_name = template_name | ||
self.status = EmailStatus.QUEUED | ||
self.created_at = datetime.now() | ||
self.updated_at = datetime.now() | ||
self.retry_count = 0 | ||
self.error = None | ||
self.scheduled_time = None | ||
|
||
class EmailQueue: | ||
def __init__(self, fastmail: FastMail, max_retries: int = 3): | ||
self.fastmail = fastmail | ||
self.max_retries = max_retries | ||
self.queue: Dict[str, QueuedEmail] = {} | ||
self._processing = False | ||
self._task = None | ||
|
||
async def add_to_queue(self, message: MessageSchema, template_name: Optional[str] = None, schedule_time: Optional[datetime] = None) -> str: | ||
queued_email = QueuedEmail(message, template_name) | ||
if schedule_time: | ||
queued_email.scheduled_time = schedule_time | ||
self.queue[queued_email.id] = queued_email | ||
return queued_email.id | ||
|
||
async def process_queue(self): | ||
self._processing = True | ||
while self._processing: | ||
now = datetime.now() | ||
for email in list(self.queue.values()): | ||
if email.status == EmailStatus.QUEUED: | ||
if email.scheduled_time and email.scheduled_time > now: | ||
continue | ||
try: | ||
email.status = EmailStatus.SENDING | ||
email.updated_at = datetime.now() | ||
await self.fastmail._send_message(email.message, email.template_name) | ||
email.status = EmailStatus.SENT | ||
email.updated_at = datetime.now() | ||
except Exception as e: | ||
email.error = str(e) | ||
email.retry_count += 1 | ||
if email.retry_count >= self.max_retries: | ||
email.status = EmailStatus.FAILED | ||
else: | ||
email.status = EmailStatus.QUEUED | ||
email.updated_at = datetime.now() | ||
await asyncio.sleep(1) | ||
|
||
def start_processing(self): | ||
if not self._task: | ||
self._task = asyncio.create_task(self.process_queue()) | ||
|
||
def stop_processing(self): | ||
self._processing = False | ||
if self._task: | ||
self._task.cancel() | ||
self._task = None | ||
|
||
def get_queue_status(self) -> Dict[str, int]: | ||
status = { | ||
"total": len(self.queue), | ||
"queued": 0, | ||
"sending": 0, | ||
"sent": 0, | ||
"failed": 0 | ||
} | ||
for email in self.queue.values(): | ||
status[email.status.value] += 1 | ||
return status | ||
|
||
def get_email_status(self, email_id: str) -> Optional[QueuedEmail]: | ||
return self.queue.get(email_id) | ||
|
||
def cancel_email(self, email_id: str) -> bool: | ||
if email_id in self.queue and self.queue[email_id].status == EmailStatus.QUEUED: | ||
del self.queue[email_id] | ||
return True | ||
return False | ||
|
||
def retry_failed(self, email_id: str) -> bool: | ||
if email_id in self.queue: | ||
email = self.queue[email_id] | ||
if email.status == EmailStatus.FAILED: | ||
email.status = EmailStatus.QUEUED | ||
email.retry_count = 0 | ||
email.error = None | ||
email.updated_at = datetime.now() | ||
return True | ||
return False |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,120 @@ | ||
import pytest | ||
from datetime import datetime, timedelta | ||
from fastapi_mail import FastMail, ConnectionConfig, MessageSchema, MessageType | ||
from fastapi_mail.queue import EmailQueue, EmailStatus | ||
import asyncio | ||
|
||
@pytest.fixture | ||
def queued_message(): | ||
return MessageSchema( | ||
subject="Test subject", | ||
recipients=["test@example.com"], | ||
body="Test email body", | ||
subtype=MessageType.plain | ||
) | ||
|
||
@pytest.fixture | ||
def queue_config(mail_config): | ||
mail_config["SUPPRESS_SEND"] = 1 | ||
return ConnectionConfig(**mail_config) | ||
|
||
@pytest.fixture | ||
def email_queue(queue_config): | ||
fastmail = FastMail(queue_config) | ||
return EmailQueue(fastmail) | ||
|
||
@pytest.mark.asyncio | ||
async def test_add_to_queue(email_queue, queued_message): | ||
queue_id = await email_queue.add_to_queue(queued_message) | ||
assert queue_id in email_queue.queue | ||
assert email_queue.queue[queue_id].status.name == "QUEUED" | ||
assert email_queue.queue[queue_id].message == queued_message | ||
assert email_queue.queue[queue_id].retry_count == 0 | ||
|
||
@pytest.mark.asyncio | ||
async def test_scheduled_email(email_queue, queued_message): | ||
schedule_time = datetime.now() + timedelta(hours=1) | ||
queue_id = await email_queue.add_to_queue(queued_message, schedule_time=schedule_time) | ||
queued_email = email_queue.queue[queue_id] | ||
assert queued_email.scheduled_time == schedule_time | ||
assert queued_email.status.name == "QUEUED" | ||
|
||
@pytest.mark.asyncio | ||
async def test_process_queue(email_queue, queued_message): | ||
queue_id = await email_queue.add_to_queue(queued_message) | ||
email_queue.start_processing() | ||
await asyncio.sleep(2) | ||
email_queue.stop_processing() | ||
processed_email = email_queue.queue[queue_id] | ||
assert processed_email.status.name == "SENT" | ||
|
||
@pytest.mark.asyncio | ||
async def test_queue_status(email_queue, queued_message): | ||
queue_id1 = await email_queue.add_to_queue(queued_message) | ||
queue_id2 = await email_queue.add_to_queue(queued_message) | ||
email_queue.queue[queue_id2].status = EmailStatus.FAILED | ||
email_queue.queue[queue_id2].retry_count = 3 | ||
status = email_queue.get_queue_status() | ||
assert status["total"] == 2 | ||
assert status["queued"] == 1 | ||
assert status["failed"] == 1 | ||
|
||
@pytest.mark.asyncio | ||
async def test_cancel_email(email_queue, queued_message): | ||
queue_id = await email_queue.add_to_queue(queued_message) | ||
success = email_queue.cancel_email(queue_id) | ||
assert success is True | ||
assert queue_id not in email_queue.queue | ||
assert email_queue.cancel_email("non-existent") is False | ||
|
||
@pytest.mark.asyncio | ||
async def test_retry_failed_email(email_queue, queued_message): | ||
queue_id = await email_queue.add_to_queue(queued_message) | ||
email_queue.queue[queue_id].status = EmailStatus.FAILED | ||
email_queue.queue[queue_id].retry_count = 3 | ||
email_queue.queue[queue_id].error = "Error occurred" | ||
success = email_queue.retry_failed(queue_id) | ||
assert success is True | ||
retried_email = email_queue.queue[queue_id] | ||
assert retried_email.status.name == "QUEUED" | ||
assert retried_email.retry_count == 0 | ||
assert retried_email.error is None | ||
|
||
@pytest.mark.asyncio | ||
async def test_get_email_status(email_queue, queued_message): | ||
queue_id = await email_queue.add_to_queue(queued_message) | ||
email_status = email_queue.get_email_status(queue_id) | ||
assert email_status is not None | ||
assert email_status.status.name == "QUEUED" | ||
assert email_queue.get_email_status("non-existent") is None | ||
|
||
@pytest.mark.asyncio | ||
async def test_max_retries(email_queue, queued_message): | ||
queue_id = await email_queue.add_to_queue(queued_message) | ||
queued_email = email_queue.queue[queue_id] | ||
for _ in range(email_queue.max_retries): | ||
queued_email.status = EmailStatus.SENDING | ||
queued_email.retry_count += 1 | ||
queued_email.status = EmailStatus.QUEUED | ||
queued_email.status = EmailStatus.SENDING | ||
queued_email.retry_count += 1 | ||
assert queued_email.retry_count > email_queue.max_retries | ||
assert queued_email.status.name == "FAILED" | ||
|
||
@pytest.mark.asyncio | ||
async def test_fastmail_queue_integration(queue_config, queued_message): | ||
fastmail = FastMail(queue_config) | ||
queue_id = await fastmail.send_message(queued_message, queue=True) | ||
assert queue_id is not None | ||
immediate_result = await fastmail.send_message(queued_message, queue=False) | ||
assert immediate_result is None | ||
|
||
@pytest.mark.asyncio | ||
async def test_scheduled_processing(email_queue, queued_message): | ||
schedule_time = datetime.now() + timedelta(seconds=2) | ||
queue_id = await email_queue.add_to_queue(queued_message, schedule_time=schedule_time) | ||
email_queue.start_processing() | ||
assert email_queue.queue[queue_id].status.name == "QUEUED" | ||
await asyncio.sleep(3) | ||
email_queue.stop_processing() | ||
assert email_queue.queue[queue_id].status.name == "SENT" |