From 69c3bac48004a4805ca47ceb934b4183aa4f0e83 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 29 Mar 2026 22:08:38 +0000
Subject: [PATCH] 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>
---
 backend/app/tests/test_gmail_client.py | 641 +++++++++++++++++++++++++
 backend/app/tests/test_mail_sources.py | 475 ++++++++++++++++++
 2 files changed, 1116 insertions(+)
 create mode 100644 backend/app/tests/test_gmail_client.py

diff --git a/backend/app/tests/test_gmail_client.py b/backend/app/tests/test_gmail_client.py
new file mode 100644
index 0000000..f7fbb66
--- /dev/null
+++ b/backend/app/tests/test_gmail_client.py
@@ -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"]
+
+
diff --git a/backend/app/tests/test_mail_sources.py b/backend/app/tests/test_mail_sources.py
index db4c95f..8efa9cb 100644
--- a/backend/app/tests/test_mail_sources.py
+++ b/backend/app/tests/test_mail_sources.py
@@ -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
 # ---------------------------------------------------------------------------