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, ) # Verify threads().get() was called with correct parameters thread_get_kwargs = ( mock_service.users.return_value.threads.return_value.get.call_args.kwargs ) assert thread_get_kwargs["userId"] == "me" assert thread_get_kwargs["id"] == "thread123" assert thread_get_kwargs["format"] == "metadata" assert "Message-ID" in thread_get_kwargs["metadataHeaders"] 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 @pytest.mark.asyncio async def test_draft_gmail_message_gracefully_degrades_when_thread_has_no_messages(): mock_service = Mock() mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"} mock_service.users().threads().get().execute.return_value = {"messages": []} 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