tests/test_gmail_handler.py

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