diff --git a/tests/gmail/test_draft_gmail_message.py b/tests/gmail/test_draft_gmail_message.py
new file mode 100644
index 0000000..b5def51
--- /dev/null
+++ b/tests/gmail/test_draft_gmail_message.py
@@ -0,0 +1,251 @@
+import base64
+import os
+import sys
+from unittest.mock import Mock
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+from core.utils import UserInputError
+from gmail.gmail_tools import draft_gmail_message
+
+
+def _unwrap(tool):
+ """Unwrap FunctionTool + decorators to the original async function."""
+ fn = tool.fn if hasattr(tool, "fn") else tool
+ while hasattr(fn, "__wrapped__"):
+ fn = fn.__wrapped__
+ return fn
+
+
+def _thread_response(*message_ids):
+ return {
+ "messages": [
+ {
+ "payload": {
+ "headers": [{"name": "Message-ID", "value": message_id}],
+ }
+ }
+ for message_id in message_ids
+ ]
+ }
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_reports_actual_attachment_count(
+ tmp_path, monkeypatch
+):
+ monkeypatch.setenv("ALLOWED_FILE_DIRS", str(tmp_path))
+ attachment_path = tmp_path / "sample.txt"
+ attachment_path.write_text("hello attachment", encoding="utf-8")
+
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft123"}
+
+ result = await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Attachment test",
+ body="Please see attached.",
+ attachments=[{"path": str(attachment_path)}],
+ include_signature=False,
+ )
+
+ assert "Draft created with 1 attachment(s)! Draft ID: draft123" in result
+
+ create_kwargs = (
+ mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+ )
+ raw_message = create_kwargs["body"]["message"]["raw"]
+ raw_bytes = base64.urlsafe_b64decode(raw_message)
+
+ assert b"Content-Disposition: attachment;" in raw_bytes
+ assert b"sample.txt" in raw_bytes
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_raises_when_no_attachments_are_added(
+ tmp_path, monkeypatch
+):
+ monkeypatch.setenv("ALLOWED_FILE_DIRS", str(tmp_path))
+ missing_path = tmp_path / "missing.txt"
+
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft123"}
+
+ with pytest.raises(UserInputError, match="No valid attachments were added"):
+ await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Attachment test",
+ body="Please see attached.",
+ attachments=[{"path": str(missing_path)}],
+ include_signature=False,
+ )
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_appends_gmail_signature_html():
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft_sig"}
+ mock_service.users().settings().sendAs().list().execute.return_value = {
+ "sendAs": [
+ {
+ "sendAsEmail": "user@example.com",
+ "isPrimary": True,
+ "signature": "
Best,
Alice
",
+ }
+ ]
+ }
+
+ result = await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Signature test",
+ body="Hello
",
+ body_format="html",
+ include_signature=True,
+ )
+
+ assert "Draft created! Draft ID: draft_sig" in result
+
+ create_kwargs = (
+ mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+ )
+ raw_message = create_kwargs["body"]["message"]["raw"]
+ raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+ assert "Hello
" in raw_text
+ assert "Best,
Alice" in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_autofills_reply_headers_from_thread():
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+ mock_service.users().threads().get().execute.return_value = _thread_response(
+ "",
+ "",
+ "",
+ )
+
+ result = await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Meeting tomorrow",
+ body="Thanks for the update.",
+ thread_id="thread123",
+ include_signature=False,
+ )
+
+ assert "Draft created! Draft ID: draft_reply" in result
+
+ create_kwargs = (
+ mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+ )
+ raw_message = create_kwargs["body"]["message"]["raw"]
+ raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+ assert "In-Reply-To: " in raw_text
+ assert (
+ "References: "
+ in raw_text
+ )
+ assert create_kwargs["body"]["message"]["threadId"] == "thread123"
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_uses_explicit_in_reply_to_when_filling_references():
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+ mock_service.users().threads().get().execute.return_value = _thread_response(
+ "",
+ "",
+ "",
+ )
+
+ await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Meeting tomorrow",
+ body="Replying to an earlier message.",
+ thread_id="thread123",
+ in_reply_to="",
+ include_signature=False,
+ )
+
+ create_kwargs = (
+ mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+ )
+ raw_message = create_kwargs["body"]["message"]["raw"]
+ raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+ assert "In-Reply-To: " in raw_text
+ assert "References: " in raw_text
+ assert "" not in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_uses_explicit_references_when_filling_in_reply_to():
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+ mock_service.users().threads().get().execute.return_value = _thread_response(
+ "",
+ "",
+ "",
+ )
+
+ await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Meeting tomorrow",
+ body="Replying to an earlier message.",
+ thread_id="thread123",
+ references=" ",
+ include_signature=False,
+ )
+
+ create_kwargs = (
+ mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+ )
+ raw_message = create_kwargs["body"]["message"]["raw"]
+ raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+ assert "In-Reply-To: " in raw_text
+ assert "References: " in raw_text
+ assert "" not in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_gracefully_degrades_when_thread_fetch_fails():
+ mock_service = Mock()
+ mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+ mock_service.users().threads().get().execute.side_effect = RuntimeError("boom")
+
+ result = await _unwrap(draft_gmail_message)(
+ service=mock_service,
+ user_google_email="user@example.com",
+ to="recipient@example.com",
+ subject="Meeting tomorrow",
+ body="Thanks for the update.",
+ thread_id="thread123",
+ include_signature=False,
+ )
+
+ assert "Draft created! Draft ID: draft_reply" in result
+
+ create_kwargs = (
+ mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+ )
+ raw_message = create_kwargs["body"]["message"]["raw"]
+ raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+ assert "In-Reply-To:" not in raw_text
+ assert "References:" not in raw_text