Spaces:
Running
Running
File size: 2,584 Bytes
43482c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | 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
|