feat: add GmailClient unit tests and main.py helper tests to meet coverage target
Agent-Logs-Url: https://github.com/christianlouis/dmarq/sessions/77ac3e98-81fa-45fa-9cc2-481896ffb48e Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
@@ -0,0 +1,641 @@
|
||||
"""
|
||||
Unit tests for app.services.gmail_client.GmailClient.
|
||||
|
||||
All external I/O (httpx, google-auth, googleapiclient) is mocked so these
|
||||
tests never make real network calls.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import email as email_mod
|
||||
import json
|
||||
from email import encoders as email_encoders
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.base import MIMEBase
|
||||
from email.mime.text import MIMEText
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.gmail_client import GmailClient, DMARC_GMAIL_QUERY
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_client(
|
||||
access_token: str = "acc",
|
||||
refresh_token: str = "ref",
|
||||
already_ingested: Optional[list] = None,
|
||||
) -> GmailClient:
|
||||
"""Instantiate a GmailClient with real Credentials mocked out."""
|
||||
with patch("app.services.gmail_client.Credentials") as mock_creds_class:
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.token = access_token
|
||||
mock_creds.refresh_token = refresh_token
|
||||
mock_creds.expired = False
|
||||
mock_creds_class.return_value = mock_creds
|
||||
client = GmailClient(
|
||||
client_id="cid",
|
||||
client_secret="csec",
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
already_ingested_ids=already_ingested or [],
|
||||
)
|
||||
# Expose the mock so tests can manipulate it
|
||||
client._mock_creds = mock_creds # type: ignore[attr-defined]
|
||||
return client
|
||||
|
||||
|
||||
def _make_raw_email(attachments: list) -> bytes:
|
||||
"""Build a raw MIME email whose attachments are described by *attachments*.
|
||||
|
||||
Each item in *attachments* is a dict with keys:
|
||||
filename: str
|
||||
content: bytes
|
||||
disposition: str (default "attachment")
|
||||
"""
|
||||
msg = MIMEMultipart()
|
||||
msg["Subject"] = "DMARC aggregate report"
|
||||
msg["From"] = "noreply@example.com"
|
||||
msg["To"] = "user@gmail.com"
|
||||
msg.attach(MIMEText("See attached DMARC report.", "plain"))
|
||||
|
||||
for att in attachments:
|
||||
part = MIMEBase("application", "octet-stream")
|
||||
part.set_payload(att["content"])
|
||||
email_encoders.encode_base64(part)
|
||||
disposition = att.get("disposition", "attachment")
|
||||
part.add_header(
|
||||
"Content-Disposition",
|
||||
disposition,
|
||||
filename=att["filename"],
|
||||
)
|
||||
msg.attach(part)
|
||||
|
||||
return msg.as_bytes()
|
||||
|
||||
|
||||
def _b64_raw(raw_bytes: bytes) -> str:
|
||||
"""URL-safe base64-encode bytes (as Gmail API returns them)."""
|
||||
return base64.urlsafe_b64encode(raw_bytes).decode()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# __init__ / basic construction
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestGmailClientInit:
|
||||
def test_init_stores_tokens(self):
|
||||
client = _make_client(access_token="my-acc", refresh_token="my-ref")
|
||||
assert client._initial_access_token == "my-acc"
|
||||
assert client.client_id == "cid"
|
||||
assert client.client_secret == "csec"
|
||||
|
||||
def test_init_already_ingested_defaults_to_empty(self):
|
||||
client = _make_client()
|
||||
assert client.already_ingested_ids == []
|
||||
|
||||
def test_init_already_ingested_is_copied(self):
|
||||
ids = ["a", "b"]
|
||||
client = _make_client(already_ingested=ids)
|
||||
assert client.already_ingested_ids == ["a", "b"]
|
||||
# Mutating the original should not affect the client
|
||||
ids.append("c")
|
||||
assert "c" not in client.already_ingested_ids
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# get_refreshed_tokens
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestGetRefreshedTokens:
|
||||
def test_returns_none_when_token_unchanged(self):
|
||||
client = _make_client(access_token="original")
|
||||
# credentials.token == _initial_access_token → no refresh happened
|
||||
client._mock_creds.token = "original"
|
||||
assert client.get_refreshed_tokens() is None
|
||||
|
||||
def test_returns_new_access_token_when_changed(self):
|
||||
client = _make_client(access_token="original")
|
||||
client._mock_creds.token = "new-token"
|
||||
client._mock_creds.refresh_token = None
|
||||
result = client.get_refreshed_tokens()
|
||||
assert result is not None
|
||||
assert result["access_token"] == "new-token"
|
||||
assert "refresh_token" not in result
|
||||
|
||||
def test_returns_both_tokens_when_refresh_token_present(self):
|
||||
client = _make_client(access_token="original")
|
||||
client._mock_creds.token = "new-token"
|
||||
client._mock_creds.refresh_token = "new-refresh"
|
||||
result = client.get_refreshed_tokens()
|
||||
assert result is not None
|
||||
assert result["access_token"] == "new-token"
|
||||
assert result["refresh_token"] == "new-refresh"
|
||||
|
||||
def test_returns_none_when_token_is_none(self):
|
||||
client = _make_client(access_token="original")
|
||||
client._mock_creds.token = None
|
||||
assert client.get_refreshed_tokens() is None
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# build_authorization_url (already partly covered; extend for completeness)
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestBuildAuthorizationUrl:
|
||||
def test_includes_all_required_params(self):
|
||||
url = GmailClient.build_authorization_url(
|
||||
client_id="cid",
|
||||
redirect_uri="https://example.com/cb",
|
||||
state="99",
|
||||
)
|
||||
assert "client_id=cid" in url
|
||||
assert "response_type=code" in url
|
||||
assert "gmail.readonly" in url
|
||||
assert "access_type=offline" in url
|
||||
assert "prompt=consent" in url
|
||||
assert "state=99" in url
|
||||
|
||||
def test_state_omitted_when_none(self):
|
||||
url = GmailClient.build_authorization_url(
|
||||
client_id="cid",
|
||||
redirect_uri="https://example.com/cb",
|
||||
)
|
||||
assert "state=" not in url
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# exchange_code_for_tokens
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestExchangeCodeForTokens:
|
||||
def test_success_returns_json(self):
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.status_code = 200
|
||||
mock_resp.json.return_value = {"access_token": "acc", "refresh_token": "ref"}
|
||||
|
||||
with patch("app.services.gmail_client.httpx.post", return_value=mock_resp):
|
||||
result = GmailClient.exchange_code_for_tokens(
|
||||
client_id="cid",
|
||||
client_secret="csec",
|
||||
code="auth-code",
|
||||
redirect_uri="https://example.com/cb",
|
||||
)
|
||||
|
||||
assert result["access_token"] == "acc"
|
||||
assert result["refresh_token"] == "ref"
|
||||
|
||||
def test_non_200_raises_value_error(self):
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.status_code = 400
|
||||
mock_resp.text = '{"error": "invalid_grant"}'
|
||||
|
||||
with patch("app.services.gmail_client.httpx.post", return_value=mock_resp):
|
||||
with pytest.raises(ValueError, match="400"):
|
||||
GmailClient.exchange_code_for_tokens(
|
||||
client_id="cid",
|
||||
client_secret="csec",
|
||||
code="bad-code",
|
||||
redirect_uri="https://example.com/cb",
|
||||
)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# get_gmail_email
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestGetGmailEmail:
|
||||
def test_returns_email_on_200(self):
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.status_code = 200
|
||||
mock_resp.json.return_value = {"email": "user@gmail.com"}
|
||||
|
||||
with patch("app.services.gmail_client.httpx.get", return_value=mock_resp):
|
||||
result = GmailClient.get_gmail_email("valid-token")
|
||||
|
||||
assert result == "user@gmail.com"
|
||||
|
||||
def test_returns_none_on_non_200(self):
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.status_code = 401
|
||||
|
||||
with patch("app.services.gmail_client.httpx.get", return_value=mock_resp):
|
||||
result = GmailClient.get_gmail_email("expired-token")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_on_exception(self):
|
||||
with patch(
|
||||
"app.services.gmail_client.httpx.get", side_effect=Exception("network error")
|
||||
):
|
||||
result = GmailClient.get_gmail_email("some-token")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _build_service
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestBuildService:
|
||||
def test_returns_service_when_not_expired(self):
|
||||
client = _make_client()
|
||||
client._mock_creds.expired = False
|
||||
|
||||
mock_service = MagicMock()
|
||||
with patch("app.services.gmail_client.build", return_value=mock_service):
|
||||
svc = client._build_service()
|
||||
|
||||
assert svc is mock_service
|
||||
|
||||
def test_refreshes_token_when_expired(self):
|
||||
client = _make_client()
|
||||
client._mock_creds.expired = True
|
||||
client._mock_creds.refresh_token = "ref"
|
||||
|
||||
mock_service = MagicMock()
|
||||
with patch("app.services.gmail_client.build", return_value=mock_service), patch(
|
||||
"app.services.gmail_client.Request"
|
||||
):
|
||||
svc = client._build_service()
|
||||
|
||||
client._mock_creds.refresh.assert_called_once()
|
||||
assert svc is mock_service
|
||||
|
||||
def test_raises_when_refresh_fails(self):
|
||||
client = _make_client()
|
||||
client._mock_creds.expired = True
|
||||
client._mock_creds.refresh_token = "ref"
|
||||
client._mock_creds.refresh.side_effect = Exception("refresh failed")
|
||||
|
||||
with patch("app.services.gmail_client.Request"), patch(
|
||||
"app.services.gmail_client.build"
|
||||
):
|
||||
with pytest.raises(Exception, match="refresh failed"):
|
||||
client._build_service()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _list_dmarc_message_ids
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestListDmarcMessageIds:
|
||||
def test_returns_empty_when_no_messages(self):
|
||||
client = _make_client()
|
||||
service = MagicMock()
|
||||
service.users.return_value.messages.return_value.list.return_value.execute.return_value = {
|
||||
"messages": []
|
||||
}
|
||||
ids = client._list_dmarc_message_ids(service)
|
||||
assert ids == []
|
||||
|
||||
def test_returns_ids_from_single_page(self):
|
||||
client = _make_client()
|
||||
service = MagicMock()
|
||||
service.users.return_value.messages.return_value.list.return_value.execute.return_value = {
|
||||
"messages": [{"id": "id1"}, {"id": "id2"}]
|
||||
}
|
||||
ids = client._list_dmarc_message_ids(service)
|
||||
assert ids == ["id1", "id2"]
|
||||
|
||||
def test_follows_next_page_token(self):
|
||||
client = _make_client()
|
||||
# First page has a nextPageToken; second page has none
|
||||
service = MagicMock()
|
||||
execute = service.users.return_value.messages.return_value.list.return_value.execute
|
||||
execute.side_effect = [
|
||||
{"messages": [{"id": "id1"}], "nextPageToken": "page2"},
|
||||
{"messages": [{"id": "id2"}]},
|
||||
]
|
||||
ids = client._list_dmarc_message_ids(service)
|
||||
assert ids == ["id1", "id2"]
|
||||
# list() should have been called twice
|
||||
assert service.users.return_value.messages.return_value.list.call_count == 2
|
||||
|
||||
def test_raises_on_http_error(self):
|
||||
from googleapiclient.errors import HttpError
|
||||
|
||||
client = _make_client()
|
||||
service = MagicMock()
|
||||
execute = service.users.return_value.messages.return_value.list.return_value.execute
|
||||
fake_error = HttpError(MagicMock(status=403), b"forbidden")
|
||||
execute.side_effect = fake_error
|
||||
with pytest.raises(HttpError):
|
||||
client._list_dmarc_message_ids(service)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _decode_part_filename
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestDecodePartFilename:
|
||||
def test_plain_ascii_filename(self):
|
||||
part = MagicMock()
|
||||
part.get_filename.return_value = "report.xml"
|
||||
assert GmailClient._decode_part_filename(part) == "report.xml"
|
||||
|
||||
def test_none_filename_returns_empty(self):
|
||||
part = MagicMock()
|
||||
part.get_filename.return_value = None
|
||||
assert GmailClient._decode_part_filename(part) == ""
|
||||
|
||||
def test_rfc2047_encoded_filename(self):
|
||||
# Build an RFC 2047 encoded filename
|
||||
encoded = "=?utf-8?b?cmVwb3J0LnhtbA==?=" # base64("report.xml")
|
||||
part = MagicMock()
|
||||
part.get_filename.return_value = encoded
|
||||
result = GmailClient._decode_part_filename(part)
|
||||
assert result == "report.xml"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _is_dmarc_attachment
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestIsDmarcAttachment:
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
[
|
||||
"report.xml",
|
||||
"report.XML", # case-insensitive
|
||||
"report.zip",
|
||||
"report.gz",
|
||||
"report.gzip",
|
||||
"Report.ZIP",
|
||||
],
|
||||
)
|
||||
def test_dmarc_extensions_return_true(self, filename):
|
||||
assert GmailClient._is_dmarc_attachment(filename) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
["report.txt", "image.png", "report.pdf", "report.tar", ""],
|
||||
)
|
||||
def test_non_dmarc_extensions_return_false(self, filename):
|
||||
assert GmailClient._is_dmarc_attachment(filename) is False
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _process_message
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestProcessMessage:
|
||||
def test_fetches_and_processes_message(self):
|
||||
"""Happy path: message fetched, attachments processed."""
|
||||
client = _make_client()
|
||||
raw_email = _make_raw_email(
|
||||
[{"filename": "report.xml", "content": b"<xml/>"}]
|
||||
)
|
||||
raw_b64 = _b64_raw(raw_email)
|
||||
|
||||
service = MagicMock()
|
||||
service.users.return_value.messages.return_value.get.return_value.execute.return_value = {
|
||||
"raw": raw_b64
|
||||
}
|
||||
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
with patch.object(client, "_process_attachments", return_value=0) as mock_proc:
|
||||
count = client._process_message(service, "msg1", stats)
|
||||
|
||||
mock_proc.assert_called_once()
|
||||
assert count == 0 # our mock returns 0
|
||||
|
||||
def test_http_error_recorded_and_returns_zero(self):
|
||||
from googleapiclient.errors import HttpError
|
||||
|
||||
client = _make_client()
|
||||
service = MagicMock()
|
||||
fake_error = HttpError(MagicMock(status=404), b"not found")
|
||||
service.users.return_value.messages.return_value.get.return_value.execute.side_effect = (
|
||||
fake_error
|
||||
)
|
||||
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
count = client._process_message(service, "bad-id", stats)
|
||||
|
||||
assert count == 0
|
||||
assert len(stats["errors"]) == 1
|
||||
assert "bad-id" in stats["errors"][0]
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _process_attachments
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestProcessAttachments:
|
||||
def test_no_attachments_returns_zero(self):
|
||||
client = _make_client()
|
||||
msg = email_mod.message_from_bytes(
|
||||
b"From: a@b.com\r\nTo: c@d.com\r\n\r\nHello"
|
||||
)
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
count = client._process_attachments(msg, stats)
|
||||
assert count == 0
|
||||
|
||||
def test_non_dmarc_attachment_skipped(self):
|
||||
"""An inline or non-DMARC file should not count as a report."""
|
||||
client = _make_client()
|
||||
raw = _make_raw_email(
|
||||
[{"filename": "photo.png", "content": b"\x89PNG"}]
|
||||
)
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
count = client._process_attachments(msg, stats)
|
||||
assert count == 0
|
||||
assert stats["reports_found"] == 0
|
||||
|
||||
def test_dmarc_xml_attachment_is_parsed(self):
|
||||
"""A .xml attachment is parsed via DMARCParser and counts as a report."""
|
||||
client = _make_client()
|
||||
raw = _make_raw_email(
|
||||
[{"filename": "report.xml", "content": b"<xml_content/>"}]
|
||||
)
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
|
||||
mock_report = {"domain": "example.com", "records": []}
|
||||
with patch("app.services.gmail_client.DMARCParser") as mock_parser_class:
|
||||
mock_parser = MagicMock()
|
||||
mock_parser.parse.return_value = [mock_report]
|
||||
mock_parser_class.return_value = mock_parser
|
||||
# Also mock report_store.add_report to avoid real persistence
|
||||
with patch.object(client.report_store, "add_report"):
|
||||
count = client._process_attachments(msg, stats)
|
||||
|
||||
assert count == 1
|
||||
assert stats["reports_found"] == 1
|
||||
|
||||
def test_dmarc_attachment_with_empty_content_skipped(self):
|
||||
"""A DMARC-named attachment with truly empty payload is skipped gracefully."""
|
||||
client = _make_client()
|
||||
# Build an attachment with empty bytes – base64 of b"" is b""
|
||||
raw = _make_raw_email(
|
||||
[{"filename": "report.zip", "content": b""}]
|
||||
)
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
count = client._process_attachments(msg, stats)
|
||||
# Empty payload → `get_payload(decode=True)` returns b"" which is
|
||||
# falsy, so the attachment is skipped
|
||||
assert count == 0
|
||||
|
||||
def test_parse_exception_adds_error_and_continues(self):
|
||||
"""A parse error should be recorded in stats but not raise."""
|
||||
client = _make_client()
|
||||
raw = _make_raw_email(
|
||||
[
|
||||
{"filename": "bad.xml", "content": b"corrupt"},
|
||||
{"filename": "good.xml", "content": b"<xml/>"},
|
||||
]
|
||||
)
|
||||
msg = email_mod.message_from_bytes(raw)
|
||||
stats = {"reports_found": 0, "errors": []}
|
||||
|
||||
good_report = {"domain": "example.com", "records": []}
|
||||
call_count = 0
|
||||
|
||||
def parse_side_effect(content, filename):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise ValueError("bad xml")
|
||||
return [good_report]
|
||||
|
||||
with patch("app.services.gmail_client.DMARCParser") as mock_parser_class:
|
||||
mock_parser = MagicMock()
|
||||
mock_parser.parse.side_effect = parse_side_effect
|
||||
mock_parser_class.return_value = mock_parser
|
||||
with patch.object(client.report_store, "add_report"):
|
||||
count = client._process_attachments(msg, stats)
|
||||
|
||||
assert len(stats["errors"]) == 1
|
||||
assert "bad.xml" in stats["errors"][0]
|
||||
assert count == 1 # second attachment still parsed
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# fetch_reports
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestFetchReports:
|
||||
def test_returns_failure_when_build_service_raises(self):
|
||||
client = _make_client()
|
||||
with patch.object(
|
||||
client, "_build_service", side_effect=Exception("auth error")
|
||||
):
|
||||
result = client.fetch_reports()
|
||||
|
||||
assert result["success"] is False
|
||||
assert "auth error" in result.get("error", "")
|
||||
|
||||
def test_returns_failure_when_list_messages_raises(self):
|
||||
client = _make_client()
|
||||
mock_service = MagicMock()
|
||||
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
|
||||
client, "_list_dmarc_message_ids", side_effect=Exception("list error")
|
||||
):
|
||||
result = client.fetch_reports()
|
||||
|
||||
assert result["success"] is False
|
||||
|
||||
def test_returns_success_with_no_messages(self):
|
||||
client = _make_client()
|
||||
mock_service = MagicMock()
|
||||
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
|
||||
client, "_list_dmarc_message_ids", return_value=[]
|
||||
):
|
||||
result = client.fetch_reports()
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["processed"] == 0
|
||||
|
||||
def test_skips_already_ingested_messages(self):
|
||||
client = _make_client(already_ingested=["id1"])
|
||||
mock_service = MagicMock()
|
||||
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
|
||||
client, "_list_dmarc_message_ids", return_value=["id1", "id2"]
|
||||
), patch.object(client, "_process_message", return_value=0) as mock_proc:
|
||||
result = client.fetch_reports()
|
||||
|
||||
# Only id2 should be processed; id1 is already ingested
|
||||
assert mock_proc.call_count == 1
|
||||
call_args = mock_proc.call_args_list[0][0]
|
||||
assert call_args[1] == "id2"
|
||||
assert result["processed"] == 1
|
||||
|
||||
def test_tracks_new_ingested_ids(self):
|
||||
client = _make_client()
|
||||
mock_service = MagicMock()
|
||||
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
|
||||
client, "_list_dmarc_message_ids", return_value=["id1", "id2"]
|
||||
), patch.object(client, "_process_message", return_value=0):
|
||||
result = client.fetch_reports()
|
||||
|
||||
assert "id1" in result["new_ingested_ids"]
|
||||
assert "id2" in result["new_ingested_ids"]
|
||||
|
||||
def test_reports_new_domains(self):
|
||||
"""fetch_reports should report domains that appear after ingestion."""
|
||||
from app.services.report_store import ReportStore
|
||||
|
||||
client = _make_client()
|
||||
mock_service = MagicMock()
|
||||
|
||||
def _process_side_effect(service, msg_id, stats):
|
||||
# Simulate adding a domain to the report store
|
||||
ReportStore.get_instance().add_report(
|
||||
{
|
||||
"org_name": "Test Org",
|
||||
"report_id": "r1",
|
||||
"begin_date": "2024-01-01",
|
||||
"end_date": "2024-01-02",
|
||||
"domain": "newdomain.example",
|
||||
"records": [],
|
||||
}
|
||||
)
|
||||
stats["reports_found"] += 1
|
||||
return 1
|
||||
|
||||
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
|
||||
client, "_list_dmarc_message_ids", return_value=["id1"]
|
||||
), patch.object(client, "_process_message", side_effect=_process_side_effect):
|
||||
result = client.fetch_reports()
|
||||
|
||||
assert "newdomain.example" in result["new_domains"]
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# load_ingested_ids / dump_ingested_ids (already tested in TestGmailClientHelpers
|
||||
# in test_mail_sources.py; add a few edge-cases here)
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestIngestedIdHelpers:
|
||||
def test_load_non_list_json_returns_no_error(self):
|
||||
# Valid JSON but not a list – should gracefully not raise
|
||||
result = GmailClient.load_ingested_ids('{"key": "value"}')
|
||||
assert result is not None # no crash
|
||||
|
||||
def test_dump_preserves_order(self):
|
||||
ids = ["z", "a", "m"]
|
||||
dumped = GmailClient.dump_ingested_ids(ids)
|
||||
loaded = json.loads(dumped)
|
||||
assert loaded == ["z", "a", "m"]
|
||||
|
||||
|
||||
@@ -1220,6 +1220,481 @@ class TestGmailTestConnectionFailure:
|
||||
assert "check server logs" in data["message"].lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for new main.py helper functions (_poll_single_gmail_source,
|
||||
# _trigger_poll_imap_source, _trigger_poll_gmail_source, _poll_source_for_trigger)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPollSingleGmailSource:
|
||||
"""Unit tests for app.main._poll_single_gmail_source."""
|
||||
|
||||
def _make_source(self, *, access_token="tok", refresh_token="ref"):
|
||||
src = MagicMock()
|
||||
src.id = 1
|
||||
src.gmail_access_token = access_token
|
||||
src.gmail_refresh_token = refresh_token
|
||||
src.gmail_client_id = "cid"
|
||||
src.gmail_client_secret = "csec"
|
||||
src.gmail_ingested_ids = "[]"
|
||||
src.gmail_email = "u@gmail.com"
|
||||
return src
|
||||
|
||||
def test_skips_when_no_access_token(self):
|
||||
"""Source without OAuth token → early return, no GmailClient created."""
|
||||
from app.main import _poll_single_gmail_source
|
||||
|
||||
src = self._make_source(access_token=None)
|
||||
|
||||
with patch("app.main.GmailClient") as mock_gc:
|
||||
_poll_single_gmail_source(src)
|
||||
|
||||
mock_gc.assert_not_called()
|
||||
|
||||
def test_fetches_reports_and_persists_ids(self):
|
||||
"""Happy-path: client is created, reports fetched, IDs saved to DB."""
|
||||
from app.main import _poll_single_gmail_source
|
||||
|
||||
src = self._make_source()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.fetch_reports.return_value = {
|
||||
"success": True,
|
||||
"processed": 2,
|
||||
"reports_found": 1,
|
||||
"new_domains": [],
|
||||
"errors": [],
|
||||
"new_ingested_ids": ["id1", "id2"],
|
||||
}
|
||||
mock_client.get_refreshed_tokens.return_value = None
|
||||
|
||||
mock_db_source = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_db.__enter__ = MagicMock(return_value=mock_db)
|
||||
mock_db.__exit__ = MagicMock(return_value=False)
|
||||
mock_db.query.return_value.get.return_value = mock_db_source
|
||||
|
||||
with patch("app.main.GmailClient", return_value=mock_client), patch(
|
||||
"app.main.SessionLocal", return_value=mock_db
|
||||
), patch("app.main.GmailClient.load_ingested_ids", return_value=[]), patch(
|
||||
"app.main.GmailClient.dump_ingested_ids", return_value='["id1","id2"]'
|
||||
):
|
||||
_poll_single_gmail_source(src)
|
||||
|
||||
mock_client.fetch_reports.assert_called_once()
|
||||
|
||||
def test_logs_new_domains_on_success(self):
|
||||
"""When results include new_domains, the function logs them."""
|
||||
from app.main import _poll_single_gmail_source
|
||||
|
||||
src = self._make_source()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.fetch_reports.return_value = {
|
||||
"success": True,
|
||||
"processed": 1,
|
||||
"reports_found": 1,
|
||||
"new_domains": ["example.com"],
|
||||
"errors": [],
|
||||
"new_ingested_ids": [],
|
||||
}
|
||||
mock_client.get_refreshed_tokens.return_value = None
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.get.return_value = MagicMock()
|
||||
|
||||
with patch("app.main.GmailClient", return_value=mock_client), patch(
|
||||
"app.main.SessionLocal", return_value=mock_db
|
||||
), patch("app.main.GmailClient.load_ingested_ids", return_value=[]):
|
||||
_poll_single_gmail_source(src) # should not raise
|
||||
|
||||
def test_logs_error_on_failure(self):
|
||||
"""When results['success'] is False, the function logs an error."""
|
||||
from app.main import _poll_single_gmail_source
|
||||
|
||||
src = self._make_source()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.fetch_reports.return_value = {
|
||||
"success": False,
|
||||
"error": "auth failed",
|
||||
"processed": 0,
|
||||
"reports_found": 0,
|
||||
"new_domains": [],
|
||||
"errors": [],
|
||||
"new_ingested_ids": [],
|
||||
}
|
||||
mock_client.get_refreshed_tokens.return_value = None
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.get.return_value = MagicMock()
|
||||
|
||||
with patch("app.main.GmailClient", return_value=mock_client), patch(
|
||||
"app.main.SessionLocal", return_value=mock_db
|
||||
), patch("app.main.GmailClient.load_ingested_ids", return_value=[]):
|
||||
_poll_single_gmail_source(src) # should not raise
|
||||
|
||||
def test_persists_refreshed_tokens(self):
|
||||
"""When GmailClient reports refreshed tokens, they are saved to the DB row."""
|
||||
from app.main import _poll_single_gmail_source
|
||||
|
||||
src = self._make_source()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.fetch_reports.return_value = {
|
||||
"success": True,
|
||||
"processed": 0,
|
||||
"reports_found": 0,
|
||||
"new_domains": [],
|
||||
"errors": [],
|
||||
"new_ingested_ids": [],
|
||||
}
|
||||
mock_client.get_refreshed_tokens.return_value = {
|
||||
"access_token": "new-acc",
|
||||
"refresh_token": "new-ref",
|
||||
}
|
||||
|
||||
mock_db_source = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.get.return_value = mock_db_source
|
||||
|
||||
with patch("app.main.GmailClient", return_value=mock_client), patch(
|
||||
"app.main.SessionLocal", return_value=mock_db
|
||||
), patch("app.main.GmailClient.load_ingested_ids", return_value=[]):
|
||||
_poll_single_gmail_source(src)
|
||||
|
||||
assert mock_db_source.gmail_access_token == "new-acc"
|
||||
assert mock_db_source.gmail_refresh_token == "new-ref"
|
||||
|
||||
|
||||
class TestTriggerPollImapSource:
|
||||
"""Unit tests for app.main._trigger_poll_imap_source."""
|
||||
|
||||
def test_returns_result_dict_on_success(self):
|
||||
from app.main import _trigger_poll_imap_source
|
||||
|
||||
src = MagicMock()
|
||||
src.id = 5
|
||||
src.name = "My IMAP"
|
||||
src.server = "imap.example.com"
|
||||
src.port = 993
|
||||
src.username = "u"
|
||||
src.password = "p"
|
||||
|
||||
mock_imap = MagicMock()
|
||||
mock_imap.fetch_reports.return_value = {
|
||||
"success": True,
|
||||
"processed": 3,
|
||||
"reports_found": 2,
|
||||
"new_domains": ["dom.example"],
|
||||
}
|
||||
|
||||
mock_db = MagicMock()
|
||||
|
||||
with patch("app.main.IMAPClient", return_value=mock_imap):
|
||||
result = _trigger_poll_imap_source(src, mock_db)
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["source_id"] == 5
|
||||
assert result["name"] == "My IMAP"
|
||||
assert result["processed"] == 3
|
||||
assert result["reports_found"] == 2
|
||||
assert result["new_domains"] == ["dom.example"]
|
||||
mock_db.commit.assert_called_once()
|
||||
|
||||
|
||||
class TestTriggerPollGmailSource:
|
||||
"""Unit tests for app.main._trigger_poll_gmail_source."""
|
||||
|
||||
def _make_src(self):
|
||||
src = MagicMock()
|
||||
src.id = 7
|
||||
src.name = "My Gmail"
|
||||
src.gmail_client_id = "cid"
|
||||
src.gmail_client_secret = "csec"
|
||||
src.gmail_access_token = "tok"
|
||||
src.gmail_refresh_token = "ref"
|
||||
src.gmail_ingested_ids = "[]"
|
||||
return src
|
||||
|
||||
def test_returns_result_dict_on_success(self):
|
||||
from app.main import _trigger_poll_gmail_source
|
||||
|
||||
src = self._make_src()
|
||||
mock_gc = MagicMock()
|
||||
mock_gc.fetch_reports.return_value = {
|
||||
"success": True,
|
||||
"processed": 1,
|
||||
"reports_found": 1,
|
||||
"new_domains": [],
|
||||
"new_ingested_ids": ["id1"],
|
||||
}
|
||||
mock_gc.get_refreshed_tokens.return_value = None
|
||||
mock_db = MagicMock()
|
||||
|
||||
with patch("app.main.GmailClient", return_value=mock_gc), patch(
|
||||
"app.main.GmailClient.load_ingested_ids", return_value=[]
|
||||
), patch("app.main.GmailClient.dump_ingested_ids", return_value='["id1"]'):
|
||||
result = _trigger_poll_gmail_source(src, mock_db)
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["source_id"] == 7
|
||||
mock_db.commit.assert_called_once()
|
||||
|
||||
def test_persists_refreshed_tokens(self):
|
||||
from app.main import _trigger_poll_gmail_source
|
||||
|
||||
src = self._make_src()
|
||||
mock_gc = MagicMock()
|
||||
mock_gc.fetch_reports.return_value = {
|
||||
"success": True,
|
||||
"processed": 0,
|
||||
"reports_found": 0,
|
||||
"new_domains": [],
|
||||
"new_ingested_ids": [],
|
||||
}
|
||||
mock_gc.get_refreshed_tokens.return_value = {
|
||||
"access_token": "new-acc",
|
||||
"refresh_token": "new-ref",
|
||||
}
|
||||
mock_db = MagicMock()
|
||||
|
||||
with patch("app.main.GmailClient", return_value=mock_gc), patch(
|
||||
"app.main.GmailClient.load_ingested_ids", return_value=[]
|
||||
):
|
||||
_trigger_poll_gmail_source(src, mock_db)
|
||||
|
||||
assert src.gmail_access_token == "new-acc"
|
||||
assert src.gmail_refresh_token == "new-ref"
|
||||
|
||||
|
||||
class TestPollSourceForTrigger:
|
||||
"""Unit tests for app.main._poll_source_for_trigger."""
|
||||
|
||||
def test_gmail_no_token_returns_skipped(self):
|
||||
from app.main import _poll_source_for_trigger
|
||||
|
||||
src = MagicMock()
|
||||
src.method = "GMAIL_API"
|
||||
src.gmail_access_token = None
|
||||
src.id = 1
|
||||
src.name = "Gmail no token"
|
||||
|
||||
result = _poll_source_for_trigger(src, MagicMock())
|
||||
|
||||
assert result["skipped"] is True
|
||||
assert "authorised" in result["reason"].lower()
|
||||
|
||||
def test_gmail_with_token_delegates_to_trigger_poll(self):
|
||||
from app.main import _poll_source_for_trigger
|
||||
|
||||
src = MagicMock()
|
||||
src.method = "GMAIL_API"
|
||||
src.gmail_access_token = "tok"
|
||||
src.id = 2
|
||||
src.name = "Gmail"
|
||||
|
||||
expected = {"source_id": 2, "name": "Gmail", "success": True}
|
||||
with patch("app.main._trigger_poll_gmail_source", return_value=expected) as mock_fn:
|
||||
result = _poll_source_for_trigger(src, MagicMock())
|
||||
|
||||
assert result is expected
|
||||
mock_fn.assert_called_once()
|
||||
|
||||
def test_gmail_exception_returns_failure_dict(self):
|
||||
from app.main import _poll_source_for_trigger
|
||||
|
||||
src = MagicMock()
|
||||
src.method = "GMAIL_API"
|
||||
src.gmail_access_token = "tok"
|
||||
src.id = 3
|
||||
src.name = "Gmail exc"
|
||||
|
||||
with patch(
|
||||
"app.main._trigger_poll_gmail_source", side_effect=Exception("boom")
|
||||
):
|
||||
result = _poll_source_for_trigger(src, MagicMock())
|
||||
|
||||
assert result["success"] is False
|
||||
assert "boom" not in result.get("error", "") # raw msg not exposed
|
||||
|
||||
def test_imap_delegates_to_trigger_poll(self):
|
||||
from app.main import _poll_source_for_trigger
|
||||
|
||||
src = MagicMock()
|
||||
src.method = "IMAP"
|
||||
src.id = 4
|
||||
src.name = "IMAP src"
|
||||
|
||||
expected = {"source_id": 4, "success": True}
|
||||
with patch("app.main._trigger_poll_imap_source", return_value=expected) as mock_fn:
|
||||
result = _poll_source_for_trigger(src, MagicMock())
|
||||
|
||||
assert result is expected
|
||||
mock_fn.assert_called_once()
|
||||
|
||||
def test_imap_exception_returns_failure_dict(self):
|
||||
from app.main import _poll_source_for_trigger
|
||||
|
||||
src = MagicMock()
|
||||
src.method = "IMAP"
|
||||
src.id = 5
|
||||
src.name = "IMAP exc"
|
||||
|
||||
with patch(
|
||||
"app.main._trigger_poll_imap_source", side_effect=Exception("imap fail")
|
||||
):
|
||||
result = _poll_source_for_trigger(src, MagicMock())
|
||||
|
||||
assert result["success"] is False
|
||||
|
||||
def test_unknown_method_returns_skipped(self):
|
||||
from app.main import _poll_source_for_trigger
|
||||
|
||||
src = MagicMock()
|
||||
src.method = "POP3"
|
||||
src.id = 6
|
||||
src.name = "POP3 src"
|
||||
|
||||
result = _poll_source_for_trigger(src, MagicMock())
|
||||
|
||||
assert result["skipped"] is True
|
||||
assert "POP3" in result["reason"]
|
||||
|
||||
|
||||
class TestPollAllEnabledSources:
|
||||
"""Unit tests for app.main._poll_all_enabled_sources dispatch logic."""
|
||||
|
||||
def test_dispatches_gmail_api_source(self):
|
||||
"""GMAIL_API sources are forwarded to _poll_single_gmail_source."""
|
||||
from app.main import _poll_all_enabled_sources
|
||||
|
||||
src = MagicMock()
|
||||
src.id = 1
|
||||
src.method = "GMAIL_API"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [src]
|
||||
|
||||
with patch("app.main.SessionLocal", return_value=mock_db), patch(
|
||||
"app.main._poll_single_gmail_source"
|
||||
) as mock_gmail:
|
||||
_poll_all_enabled_sources()
|
||||
|
||||
mock_gmail.assert_called_once_with(src)
|
||||
|
||||
def test_dispatches_imap_source(self):
|
||||
"""IMAP sources are forwarded to _poll_single_imap_source."""
|
||||
from app.main import _poll_all_enabled_sources
|
||||
|
||||
src = MagicMock()
|
||||
src.id = 2
|
||||
src.method = "IMAP"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [src]
|
||||
|
||||
with patch("app.main.SessionLocal", return_value=mock_db), patch(
|
||||
"app.main._poll_single_imap_source"
|
||||
) as mock_imap:
|
||||
_poll_all_enabled_sources()
|
||||
|
||||
mock_imap.assert_called_once_with(src)
|
||||
|
||||
def test_gmail_exception_is_caught(self):
|
||||
"""Exception from _poll_single_gmail_source must not propagate."""
|
||||
from app.main import _poll_all_enabled_sources
|
||||
|
||||
src = MagicMock()
|
||||
src.id = 3
|
||||
src.method = "GMAIL_API"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [src]
|
||||
|
||||
with patch("app.main.SessionLocal", return_value=mock_db), patch(
|
||||
"app.main._poll_single_gmail_source", side_effect=Exception("crash")
|
||||
):
|
||||
_poll_all_enabled_sources() # should not raise
|
||||
|
||||
def test_imap_exception_is_caught(self):
|
||||
"""Exception from _poll_single_imap_source must not propagate."""
|
||||
from app.main import _poll_all_enabled_sources
|
||||
|
||||
src = MagicMock()
|
||||
src.id = 4
|
||||
src.method = "IMAP"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [src]
|
||||
|
||||
with patch("app.main.SessionLocal", return_value=mock_db), patch(
|
||||
"app.main._poll_single_imap_source", side_effect=Exception("imap crash")
|
||||
):
|
||||
_poll_all_enabled_sources() # should not raise
|
||||
|
||||
def test_unknown_method_skipped(self):
|
||||
"""An unknown method logs a skip message and does not raise."""
|
||||
from app.main import _poll_all_enabled_sources
|
||||
|
||||
src = MagicMock()
|
||||
src.id = 5
|
||||
src.method = "POP3"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [src]
|
||||
|
||||
with patch("app.main.SessionLocal", return_value=mock_db):
|
||||
_poll_all_enabled_sources() # should not raise
|
||||
|
||||
|
||||
class TestTriggerPollEndpoint:
|
||||
"""Tests for the POST /api/v1/admin/trigger-poll endpoint with sources."""
|
||||
|
||||
def test_trigger_poll_with_enabled_sources(self):
|
||||
"""With enabled sources, the endpoint dispatches and returns results."""
|
||||
from app.main import app as main_app
|
||||
from app.core.security import require_admin_auth
|
||||
|
||||
async def mock_auth():
|
||||
return {"auth_type": "api_key"}
|
||||
|
||||
main_app.dependency_overrides[require_admin_auth] = mock_auth
|
||||
|
||||
try:
|
||||
mock_source = MagicMock()
|
||||
mock_source.id = 1
|
||||
mock_source.name = "Trigger GMAIL"
|
||||
mock_source.method = "GMAIL_API"
|
||||
mock_source.enabled = True
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.filter.return_value.all.return_value = [mock_source]
|
||||
|
||||
mock_result = {
|
||||
"source_id": 1,
|
||||
"name": "Trigger GMAIL",
|
||||
"success": True,
|
||||
"processed": 0,
|
||||
"reports_found": 0,
|
||||
"new_domains": [],
|
||||
}
|
||||
|
||||
with TestClient(main_app) as tc:
|
||||
with patch("app.main.SessionLocal", return_value=mock_db), patch(
|
||||
"app.main._poll_source_for_trigger", return_value=mock_result
|
||||
):
|
||||
resp = tc.post("/api/v1/admin/trigger-poll")
|
||||
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert "sources" in data
|
||||
assert len(data["sources"]) == 1
|
||||
assert data["sources"][0]["success"] is True
|
||||
finally:
|
||||
main_app.dependency_overrides.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest marker to avoid warnings for test methods without assertions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user