import pytest from unittest.mock import patch, MagicMock import pandas as pd from connectors.gmail_handler import GmailHandler @pytest.fixture def gmail_handler(): return GmailHandler() @patch('connectors.gmail_handler.build') @patch('connectors.gmail_handler.service_account.Credentials.from_service_account_info') def test_fetch_emails_service_account(mock_creds_info, mock_build, gmail_handler): mock_service = MagicMock() mock_build.return_value = mock_service # Setup mock for messages().list() mock_list = MagicMock() mock_list.execute.return_value = { "messages": [{"id": "msg1"}, {"id": "msg2"}] } mock_service.users().messages().list.return_value = mock_list # Setup mock for messages().get() mock_get1 = MagicMock() mock_get1.execute.return_value = { "payload": { "headers": [ {"name": "From", "value": "alice@example.com"}, {"name": "Subject", "value": "Hello"} ], "body": {"data": "d29ybGQ="} # 'world' in base64 }, "snippet": "world" } mock_get2 = MagicMock() mock_get2.execute.return_value = { "payload": { "headers": [ {"name": "From", "value": "bob@example.com"}, {"name": "Subject", "value": "Hi"} ], "body": {"size": 0} # No data, rely on snippet }, "snippet": "snippet text" } mock_service.users().messages().get.side_effect = [mock_get1, mock_get2] df = gmail_handler.fetch_emails({"project_id": "test"}, auth_type="service_account", num_emails=2) assert not df.empty assert len(df) == 2 assert df.iloc[0]["Sender"] == "alice@example.com" assert df.iloc[0]["Subject"] == "Hello" assert df.iloc[0]["Content"] == "Subject: Hello\n\nworld" assert df.iloc[1]["Sender"] == "bob@example.com" assert df.iloc[1]["Subject"] == "Hi" assert df.iloc[1]["Content"] == "Subject: Hi\n\nsnippet text" @patch('connectors.gmail_handler.build') @patch('connectors.gmail_handler.google.oauth2.credentials.Credentials') def test_fetch_emails_oauth2(mock_creds_cls, mock_build, gmail_handler): mock_service = MagicMock() mock_build.return_value = mock_service # Setup mock for messages().list() empty mock_list = MagicMock() mock_list.execute.return_value = {} mock_service.users().messages().list.return_value = mock_list df = gmail_handler.fetch_emails({"access_token": "token"}, auth_type="oauth2_token", num_emails=2) assert df.empty