From a87ac1737d44484a41d564c348598f367d694e49 Mon Sep 17 00:00:00 2001 From: Taylor Wilsdon Date: Tue, 17 Mar 2026 14:34:21 -0400 Subject: [PATCH] add test --- tests/gmail/test_draft_gmail_message.py | 251 ++++++++++++++++++++++++ 1 file changed, 251 insertions(+) create mode 100644 tests/gmail/test_draft_gmail_message.py 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